code: def _test_equivalence_to_reference_code(
      self, model_cls, inp, input_signature, common_kwargs, *test_kwargs
  ):
      ref_model = model_cls(use_reference_code=True, **common_kwargs)
      rng = fastmath.random.get_prng(123)
      weights, state = ref_model.init(input_signature, rng)
  
      ref_all = self._run_forward_and_backward(
          ref_model, inp, weights, state
      )
      ref_out, ref_state, ref_inp_grad, ref_weights_grad = ref_all
  
      for kwargs in test_kwargs:
          test_model = model_cls(**common_kwargs, **kwargs)
          state = test_model.init(input_signature, rng)[1]
          test_all = self._run_forward_and_backward(
              test_model, inp, weights, state
          )
          (
              test_out,
              test_state,
              test_inp_grad,
              test_weights_grad,
          ) = test_all
  
          self.assertEqual(
              jax.tree_structure(ref_out), jax.tree_structure(test_out)
          )
          self.assertEqual(
              jax.tree_structure(ref_state),
              jax.tree_structure(test_state),
          )
          self.assertEqual(
              jax.tree_structure(ref_inp_grad),
              jax.tree_structure(test_inp_grad),
          )
          self.assertEqual(
              jax.tree_structure(ref_weights_grad),
              jax.tree_structure(test_weights_grad),
          )
  
          check_close = lambda x, y: self.assertAllClose(
              x, y, rtol=2e-3, atol=2e-3
          )
          fastmath.nested_map_multiarg(check_close, ref_out, test_out)
          fastmath.nested_map_multiarg(
              check_close, ref_state, test_state
          )
          fastmath.nested_map_multiarg(
              check_close, ref_inp_grad, test_inp_grad
          )
          fastmath.nested_map_multiarg(
              check_close, ref_weights_grad, test_weights_grad
          )
  
imports: from absl.testing import parameterized
  import jax
  import numpy
  from tensorflow import test
  from trax import fastmath
  from trax import shapes
  from trax.fastmath import numpy
  from trax.layers.research import efficient_attention
label_desc: Use image node_id 6 for calling the EfficientAttentionTest obj's underlying member method code with example usage: obj._test_equivalence_to_reference_code(model_cls, inp, input_signature, common_kwargs) without return types