title: Max function in PyTorch versus in NumPy
body: <p>I am developing an algorithm that involves a CPU case, allowing for NumPy, and a GPU case, allowing for PyTorch. The object will almost always be 4D. Two versions of the object are as follows.</p>
  <pre><code>B = [
      [[[0.5000, 0.5625],
        [0.5000, 0.5625]],
  
       [[1.2500, 0.5000],
        [0.5625, 0.6875]],
  
       [[0.5625, 0.6250],
        [0.5000, 0.5625]]]
  ] 
  
  B_array = np.array(B)
  
  B_tensor = torch.Tensor(B)
  </code></pre>
  <p>I want to take the max of each 2D matrix, such that I get a result of:</p>
  <pre><code>max_array_fn(B_array)   # returns array([0.5625, 1.250, 0.6250])
  max_tensor_fn(B_tensor) # returns tensor([0.5625, 1.250, 0.6250])
  </code></pre>
  <p>Part of the solution was discussed here, but this is only for NumPy on CPU:</p>
  <p><a href="https://stackoverflow.com/questions/63368978/max-of-each-2d-matrix-in-4d-numpy-array">Max of each 2D matrix in 4D NumPy array</a></p>
  <p>However, on the GPU, it seems that PyTorch does not use the same convention as NumPy.</p>
  <p>If this was defined as a NumPy array, we could solve the problem using <code>np.array(B, axis=(0,2,3))</code>. The usage of axis in PyTorch similar to the NumPy example is not supported, as suggested here:</p>
  <p><a href="https://stackoverflow.com/questions/61156894/pytorch-torch-max-over-multiple-dimensions">PyTorch torch.max over multiple dimensions</a></p>
  <p>Is there an alternative vectorization of the method? Why is it only able to be vectorized on the CPU with NumPy and not on GPU with PyTorch?</p>
  <p>The correct solution should not use any loops and ideally one function call.</p>