code: def test_trainvar_and_ref_assign(self):
      m = objax.ModuleList([objax.TrainVar(jn.zeros(2))])
      m.append(objax.TrainRef(m[0]))
  
      def increase(x):
          m[0].assign(m[0].value + 1)
          m[1].assign(m[1].value + 1)
          return x + 1
  
      x = np.arange(10)[:, None]
      vec_increase = objax.Vectorize(increase, m.vars())
      y = vec_increase(x)
      self.assertEqual(y.tolist(), np.arange(1, 11)[:, None].tolist())
      self.assertEqual(m[0].value.tolist(), [2.0, 2.0])
  
file_path: objax/tests/vectorize.py
imports: import unittest
  import jax.numpy
  import numpy
  import objax
label_desc: Use image node_id 12 for calling the TestVectorize obj's underlying member method code with example usage: obj.test_trainvar_and_ref_assign() and returns: unknown