class_docstr: Test if PrivateGradValues group data into microbatches correctly.
class_docstr_tok: ["Test","if","PrivateGradValues","group","data","into","microbatches","correctly","."]
docstr: Test for batch_axis = (0,) with different microbatch.
docstr_tok: ["Test","for","batch_axis","=","(","0",",",")","with","different","microbatch","."]
code: def test_private_gradvalues_axis0(self):
      """Test for batch_axis = (0,) with different microbatch."""
      for use_norm_accumulation in [True, False]:
          # microbatch = 1
          g_priv = self.get_g(
              microbatch=1,
              l2_norm_clip=1.0,
              use_norm_accumulation=use_norm_accumulation,
          )
          grad_expected = -np.array(
              [1, 1, 1, -1, 1, -1], dtype=np.float32
          ).mean()
          np.testing.assert_allclose(
              g_priv[0],
              grad_expected,
              atol=1e-5,
              err_msg="microbatch=1",
          )
  
          # microbatch = 6
          g_priv = self.get_g(
              microbatch=6,
              l2_norm_clip=1.0,
              use_norm_accumulation=use_norm_accumulation,
          )
          grad_expected = 0.0
          np.testing.assert_allclose(
              g_priv[0],
              grad_expected,
              atol=1e-5,
              err_msg="microbatch=6",
          )
  
          # microbatch = 2
          g_priv = self.get_g(
              microbatch=2,
              l2_norm_clip=3.0,
              use_norm_accumulation=use_norm_accumulation,
          )
          grad_expected = -np.array(
              [2.5, -3, 2], dtype=np.float32
          ).mean()  # unclipped gradient is [2.5, -4.5, 2]
          np.testing.assert_allclose(
              g_priv[0],
              grad_expected,
              atol=1e-5,
              err_msg="microbatch=2",
          )
  
imports: import inspect
  import unittest
  from typing import Tuple, Dict, List
  import jax
  import jax.numpy
  import numpy
  from scipy.stats.distributions import chi2
  import objax
  from objax.typing import JaxArray
  from objax.zoo.dnnet import DNNet
label_desc: Use image node_id 2 for calling the TestPrivateGradValuesAxis obj's underlying member method code with example usage: obj.test_private_gradvalues_axis0() without return types