docstr: JIT parallel reductions.
docstr_tok: ["JIT","parallel","reductions","."]
code: def test_jit_parallel_reducers(self):
      """JIT parallel reductions."""
      f = objax.nn.Linear(3, 4)
      x = objax.random.normal((96, 3))
      y = f(x)
      zl = []
      for reduce in (
          lambda x: x,
          lambda x: x[0],
          lambda x: x.mean(0),
          lambda x: x.sum(0),
      ):
          fp = objax.Jit(objax.Parallel(f, reduce=reduce))
          with fp.vars().replicate():
              zl.append(fp(x))
      znone, zfirst, zmean, zsum = zl
      self.assertAlmostEqual(
          jn.square(jn.array(jn.split(y, 8)) - znone).sum(), 0, places=8
      )
      self.assertAlmostEqual(
          jn.square(jn.split(y, 8)[0] - zfirst).sum(), 0, places=8
      )
      self.assertAlmostEqual(
          jn.square(np.mean(jn.split(y, 8), 0) - zmean).sum(),
          0,
          places=8,
      )
      self.assertAlmostEqual(
          jn.square(np.sum(jn.split(y, 8), 0) - zsum).sum(), 0, places=8
      )
  
imports: import os
  import unittest
  import jax.numpy
  import numpy
  import objax
label_desc: Use image node_id 14 for calling the TestParallel obj's underlying member method code with example usage: obj.test_jit_parallel_reducers() without return types