docstr: Parallel train op.
docstr_tok: ["Parallel","train","op","."]
code: def test_parallel_train_op(self):
      """Parallel train op."""
      f = objax.nn.Sequential(
          [
              objax.nn.Linear(3, 4),
              objax.functional.relu,
              objax.nn.Linear(4, 2),
          ]
      )
      centers = objax.random.normal((1, 2, 3))
      centers *= objax.functional.rsqrt(
          (centers**2).sum(2, keepdims=True)
      )
      x = (
          objax.random.normal((256, 1, 3), stddev=0.1) + centers
      ).reshape((512, 3))
      y = jn.concatenate(
          [
              jn.zeros((256, 1), dtype=jn.uint32),
              jn.ones((256, 1), dtype=jn.uint32),
          ],
          axis=1,
      )
      y = y.reshape((512,))
      opt = objax.optimizer.Momentum(f.vars())
      all_vars = f.vars("f") + opt.vars("opt")
  
      def loss(x, y):
          xe = objax.functional.loss.cross_entropy_logits_sparse(
              f(x), y
          )
          return xe.mean()
  
      gv = objax.GradValues(loss, f.vars())
  
      def train_op(x, y):
          g, v = gv(x, y)
          opt(0.05, g)
          return v
  
      tensors = all_vars.tensors()
      loss_value = np.array([train_op(x, y)[0] for _ in range(10)])
      var_values = {k: v.value for k, v in all_vars.items()}
      all_vars.assign(tensors)
  
      self.assertGreater(loss_value.min(), 0)
  
      def train_op_para(x, y):
          g, v = gv(x, y)
          opt(0.05, objax.functional.parallel.pmean(g))
          return objax.functional.parallel.pmean(v)
  
      fp = objax.Parallel(
          train_op_para, vc=all_vars, reduce=lambda x: x[0]
      )
      with all_vars.replicate():
          loss_value_p = np.array([fp(x, y)[0] for _ in range(10)])
      var_values_p = {k: v.value for k, v in all_vars.items()}
  
      self.assertLess(jn.abs(loss_value_p / loss_value - 1).max(), 1e-6)
      for k, v in var_values.items():
          self.assertLess(
              ((v - var_values_p[k]) ** 2).sum(), 1e-12, msg=k
          )
  
imports: import os
  import unittest
  import jax.numpy
  import numpy
  import objax
label_desc: Use image node_id 8 for calling the TestParallel obj's underlying member method code with example usage: obj.test_parallel_train_op() and returns: xe, v, objax