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