I’m not sure if your theoretical analysis is correct, as it seems you are comparing a nested loop approach while the native implementation would use the TensorIterator with the at::parallel_for approach.
You could check the source code via max_out, which should call into the max_stub,  max_kernel_impl, compare_base_kernel_core 1, and eventually into TensorIteratorBase::for_each.