code: def _test_primitive(
      self, primitive: Optional[Primitive], shapes, dtype, params
  ):
      xs = _get_inputs(shapes, dtype)
      n = len(xs)
      eqn, f = _get_f_and_eqn(params, primitive, *xs)
  
      out = f(*xs)
      cts_in = ShapedArray(out.shape, out.dtype)
  
      argnums = tuple(range(n))
      js_fwd = jax.jacfwd(f, argnums)(*xs)
      js_rev = jax.jacrev(f, argnums)(*xs)
  
      for idx in range(n):
          if primitive == lax.conv_general_dilated_p and idx == 0:
              raise absltest.SkipTest(
                  "Jacobian of CNN wrt inputs not implemented."
              )
  
          if primitive == lax.div_p and idx == 1:
              raise absltest.SkipTest(
                  "Division is linear only in the first arg."
              )
  
          invals = _get_invals(idx, *xs)
          j_fwd, j_rev = js_fwd[idx], js_rev[idx]
  
          if primitive in rules.JACOBIAN_RULES:
              j_rule = rules.JACOBIAN_RULES[primitive](
                  eqn, idx, invals, cts_in
              )
          else:
              warnings.warn(
                  f"Jacobian rule for {primitive} at position {idx} not "
                  f"found."
              )
              j_rule = None
  
          with self.subTest(f"Jacobian ({idx})"):
              self._compare_jacobians(j_fwd, j_rev, j_rule, primitive)
  
          structure = rules.STRUCTURE_RULES[primitive](
              eqn, idx, invals, cts_in
          )
  
          j = j_fwd if j_rule is None else j_rule
  
          if primitive == lax.reshape_p:
              out_ndim = xs[0].ndim
              j = j.transpose(
                  tuple(
                      xs[0].ndim + i
                      for i in np.argsort(structure.in_trace)
                  )
                  + tuple(i for i in np.argsort(structure.in_trace))
              )
              j = j.reshape(
                  xs[0].shape
                  + tuple(
                      xs[0].shape[i]
                      for i in np.argsort(structure.in_trace)
                  )
              )
  
          else:
              out_ndim = out.ndim
  
          with self.subTest(f"Diagonal axes ({idx})"):
              for i, o in zip(
                  structure.in_diagonal, structure.out_diagonal
              ):
                  self._assert_is_diagonal(
                      j=j,
                      axis1=out_ndim + i[idx],
                      axis2=o,
                      constant_diagonal=False,
                  )
  
          with self.subTest(f"Constant diagonal axes ({idx})"):
              for i, o in zip(structure.in_trace, structure.out_trace):
                  self._assert_is_diagonal(
                      j=j,
                      axis1=out_ndim + i,
                      axis2=o,
                      constant_diagonal=True,
                  )
  
          with self.subTest(f"Input broadcast axes ({idx})"):
              for i in structure.in_broadcast:
                  self._assert_constant(j=j, axis=i)
  
          with self.subTest(f"Output broadcast axes ({idx})"):
              for i in structure.out_broadcast:
                  self._assert_constant(j=j, axis=i)
  
imports: import itertools
  import logging
  import random
  from typing import Optional, Sequence
  import warnings
  from absl.testing import absltest
  import jax
  from jax import lax
  from jax import config
  from jax.core import Primitive
  from jax.core import ShapedArray
  from jax.interpreters import ad
  import jax.numpy
  import more_itertools
  from neural_tangents._src.utils import rules
  from tests import test_utils
  import numpy
label_desc: Use image node_id 4 for calling the JacobianRulesTest obj's underlying member method code with example usage: obj._test_primitive(primitive, shapes, dtype, params) without return types