docstr: Test mean squared error on x and y with keep_axis.
docstr_tok: ["Test","mean","squared","error","on","x","and","y","with","keep_axis","."]
code: def test_on_mse(self):
      """Test mean squared error on x and y with keep_axis."""
      x = jn.array(
          [
              [0.76, 0.38, 0.51, 0.59],
              [0.89, 0.76, 0.17, 0.19],
              [0.14, 0.29, 0.49, 0.55],
          ]
      )
      y = jn.array(
          [[1, 0, 4, 3], [4, 0, 4, 3], [0, 0, 0, 2]], dtype=jn.uint32
      )
      e = objax.functional.loss.mean_squared_error(x, y, None)
      gold = jn.array(4.454267)
      self.assertAlmostEqual(
          jn.abs(e - gold).sum(), 0, delta=ERROR_TOLERANCE
      )
      e = objax.functional.loss.mean_squared_error(x, y)
      gold = jn.array([4.54755, 8.203675, 0.61157507])
      self.assertAlmostEqual(
          jn.abs(e - gold).sum(), 0, delta=ERROR_TOLERANCE
      )
      e = objax.functional.loss.mean_squared_error(x, y, (1,))
      gold = jn.array([3.249767, 0.2687, 9.0297, 5.2689])
      self.assertAlmostEqual(
          jn.abs(e - gold).sum(), 0, delta=ERROR_TOLERANCE
      )
  
file_path: objax/tests/loss.py
imports: import unittest
  import jax.numpy
  import objax
label_desc: Use image node_id 2 for calling the TestLoss obj's underlying member method code with example usage: obj.test_on_mse() without return types