class_docstr: This tests other `Op`\s to ensure they keep the dimensions of their inputs correctly.
class_docstr_tok: ["This","tests","other","`","Op","`","\\s","to","ensure","they","keep","the","dimensions","of","their","inputs","correctly","."]
code: def test_norm(self, axis):
      x = dtensor3()
      a = np.random.random((3, 2, 4)).astype(pytensor.config.floatX)
      mode = Mode(optimizer="fast_compile", linker="py")
  
      f = function(
          [x],
          [
              x.norm(L=1, axis=axis, keepdims=True),
              self.makeKeepDims_local(
                  x, x.norm(L=1, axis=axis, keepdims=False), axis
              ),
          ],
          mode=mode,
      )
  
      ans1, ans2 = f(a)
      assert np.allclose(ans1, ans2)
      assert ans1.shape == ans2.shape
  
      g = function(
          [x],
          [
              x.norm(L=2, axis=axis, keepdims=True),
              self.makeKeepDims_local(
                  x, x.norm(L=2, axis=axis, keepdims=False), axis
              ),
          ],
          mode=mode,
      )
  
      ans1, ans2 = g(a)
      assert np.allclose(ans1, ans2)
      assert ans1.shape == ans2.shape
  
file_path: pytensor/tests/tensor/test_keepdims.py
imports: import numpy
  import pytest
  import pytensor
  from pytensor import function
  from pytensor.compile.mode import Mode
  from pytensor.tensor.elemwise import DimShuffle
  from pytensor.tensor.math import all
  from pytensor.tensor.math import any
  from pytensor.tensor.math import argmax, argmin, max_and_argmax, mean, prod, std, var
  from pytensor.tensor.math import max
  from pytensor.tensor.math import min
  from pytensor.tensor.math import sum
  from pytensor.tensor.type import dtensor3
label_desc: Use image node_id 5 for calling the TestKeepDims obj's underlying member method code with example usage: obj.test_norm(axis) without return types