docstr: Parallel weight decay.
docstr_tok: ["Parallel","weight","decay","."]
code: def test_parallel_weight_decay(self):
      """Parallel weight decay."""
      f = objax.nn.Sequential(
          [objax.nn.Linear(3, 4), objax.nn.Linear(4, 2)]
      )
      fvars = f.vars()
  
      def loss_fn():
          return 0.5 * sum(
              (v.value**2).sum()
              for k, v in fvars.items()
              if k.endswith(".w")
          )
  
      tensors = fvars.tensors()
      loss_value = loss_fn()
      fvars.assign(tensors)
      self.assertGreater(loss_value, 0)
  
      fp = objax.Parallel(loss_fn, vc=fvars, reduce=lambda x: x[0])
      with fvars.replicate():
          loss_value_p = fp()
      self.assertLess(abs(loss_value_p / loss_value - 1), 1e-6)
  
imports: import os
  import unittest
  import jax.numpy
  import numpy
  import objax
label_desc: Use image node_id 9 for calling the TestParallel obj's underlying member method code with example usage: obj.test_parallel_weight_decay() and returns: unknown