docstr: Parallel EMA with weight sharing.
docstr_tok: ["Parallel","EMA","with","weight","sharing","."]
code: def test_parallel_ema_shared(self):
      """Parallel EMA with weight sharing."""
      f = objax.nn.Sequential(
          [
              objax.nn.Linear(3, 4),
              objax.nn.BatchNorm0D(4),
              objax.nn.Linear(4, 2),
              objax.nn.Dropout(1),
          ]
      )
      ema = objax.optimizer.ExponentialMovingAverage(
          f.vars().subset(objax.TrainVar)
      )
      ema_f = ema.replace_vars(f)
      ema()
      all_vars = f.vars() + ema.vars()
  
      ema_fp = objax.Parallel(
          lambda x: ema_f(x, training=False), all_vars
      )
      ema_fps = objax.Parallel(
          lambda x: ema_f(x, training=False),
          all_vars + f.vars("shared"),
      )
      x = objax.random.normal((96, 3))
      with all_vars.replicate():
          z = ema_fp(x)
          zs = ema_fps(x)
      self.assertLess(jn.abs(z - zs).mean(), 1e-6)
  
imports: import os
  import unittest
  import jax.numpy
  import numpy
  import objax
label_desc: Use image node_id 11 for calling the TestParallel obj's underlying member method code with example usage: obj.test_parallel_ema_shared() without return types