docstr: Test TrainRef behavior.
docstr_tok: ["Test","TrainRef","behavior","."]
code: def test_state_ref(self):
      """Test TrainRef behavior."""
      tv = objax.TrainVar(
          jn.arange(6, dtype=jn.float32).reshape((2, 3))
      )
      v = objax.TrainRef(tv)
      for vi in (v, tv):
          self.assertEqual(vi.value.shape, (2, 3))
          self.assertEqual(vi.value.sum(), 15)
      v.value += 2
      self.assertEqual(v.value.sum(), 27)
      self.assertEqual(tv.value.sum(), 27)
      v.assign(v.value - 1)
      self.assertEqual(v.value.sum(), 21)
      self.assertEqual(tv.value.sum(), 21)
      v.reduce(v.value)
      for vi in (v, tv):
          self.assertEqual(vi.value.shape, (2, 3))
          self.assertEqual(vi.value.tolist(), [[1, 2, 3], [4, 5, 6]])
      tv.reduce(tv.value)
      for vi in (v, tv):
          self.assertEqual(vi.value.shape, (3,))
          self.assertEqual(vi.value.tolist(), [2.5, 3.5, 4.5])
  
imports: import unittest
  import numpy
  import jax.numpy
  import objax
label_desc: Use image node_id 3 for calling the TestVariable obj's underlying member method code with example usage: obj.test_state_ref() without return types