docstr: JIT parallel inference (concat reduction) with batch norm in train mode.
docstr_tok: ["JIT","parallel","inference","(","concat","reduction",")","with","batch","norm","in","train","mode","."]
code: def test_jit_parallel_bntrain_concat(self):
      """JIT parallel inference (concat reduction) with batch norm in train mode."""
      f = objax.nn.Sequential(
          [
              objax.nn.Conv2D(3, 4, k=1),
              objax.nn.BatchNorm2D(4),
              objax.nn.Conv2D(4, 2, k=1),
          ]
      )
      x = objax.random.normal((64, 3, 16, 16))
      states, values = [], []
      tensors = f.vars().tensors()
      for it in range(2):
          values.append(f(x, training=True))
          states += [f[1].running_var.value, f[1].running_mean.value]
      f.vars().assign(tensors)
  
      self.assertAlmostEqual(
          ((values[0] - values[1]) ** 2).sum(), 0, places=8
      )
      self.assertGreater(((states[0] - states[2]) ** 2).sum(), 0)
      self.assertGreater(((states[1] - states[3]) ** 2).sum(), 0)
  
      fp = objax.Jit(
          objax.Parallel(lambda x: f(x, training=True), vc=f.vars())
      )
      x8 = jn.broadcast_to(x, (8, 64, 3, 16, 16)).reshape(
          (-1, 3, 16, 16)
      )
      tensors = fp.vars().tensors()
      for it in range(2):
          with fp.vars().replicate():
              z = fp(x8).reshape((-1,) + values[it].shape)
          self.assertAlmostEqual(
              ((f[1].running_var.value - states[2 * it]) ** 2).sum(),
              0,
              delta=1e-12,
          )
          self.assertAlmostEqual(
              (
                  (f[1].running_mean.value - states[2 * it + 1]) ** 2
              ).sum(),
              0,
              delta=1e-12,
          )
          self.assertLess(((z - values[it][None]) ** 2).sum(), 1e-6)
      fp.vars().assign(tensors)
  
imports: import os
  import unittest
  import jax.numpy
  import numpy
  import objax
label_desc: Use image node_id 13 for calling the TestParallel obj's underlying member method code with example usage: obj.test_jit_parallel_bntrain_concat() without return types