code: def test_fp8_dot_general_injection(self):
      # Used to cast the inputs to be representable in FP8, so that the difference
      # of the results from the original gemm and fp8 gemm is small.
      cast_to_representable = functools.partial(
          fp8_ops.quantize_dequantize,
          scale=jnp.ones((1,)),
          compute_dtype=jnp.float32,
      )
  
      init_key, random_key = random.split(random.PRNGKey(seed=123), 2)
      x = cast_to_representable(
          random.uniform(random_key, (16, 32)), jnp.float8_e4m3fn
      )
      dy = cast_to_representable(
          random.uniform(random_key, (16, 64)), jnp.float8_e5m2
      )
  
      def run(fp8_injection, expected_shapes):
          p = nn.DenseGeneral(features=64, name="dense")
          if fp8_injection:
              p.dot_general_cls = nn.Fp8DotGeneralOp
          y, initial_vars = p.init_with_output(init_key, x)
          var_shapes = jax.tree_util.tree_map(jnp.shape, initial_vars)
          self.assertEqual(var_shapes, expected_shapes)
  
          def _train(variables, x):
              y = p.apply(variables, x)
              loss = y * dy
              return jnp.mean(loss)
  
          train_fn = jax.jit(jax.value_and_grad(_train, argnums=[0, 1]))
          outputs, grads = train_fn(initial_vars, x)
          return outputs, grads
  
      expected_shapes_original = {
          "params": {"kernel": (32, 64), "bias": (64,)},
      }
      expected_shapes_new = {
          "params": {"kernel": (32, 64), "bias": (64,)},
          fp8_ops.OVERWRITE_WITH_GRADIENT: {
              "Fp8DotGeneralOp_0": {
                  "input_amax_history": (1024,),
                  "kernel_amax_history": (1024,),
                  "output_grad_amax_history": (1024,),
                  "input_scale": (1,),
                  "kernel_scale": (1,),
                  "output_grad_scale": (1,),
              }
          },
      }
  
      output1a, output1b = run(False, expected_shapes_original)
      output2a, output2b = run(True, expected_shapes_new)
      dw1, dw2 = (
          output1b[0]["params"]["kernel"],
          output2b[0]["params"]["kernel"],
      )
      dx1, dx2 = output1b[1], output2b[1]
  
      np.testing.assert_allclose(output1a, output2a, atol=1e-02)
      np.testing.assert_allclose(dw1, dw2, atol=1e-04)
      np.testing.assert_allclose(dx1, dx2, atol=1e-04)
  
imports: import copy
  import functools
  from typing import Any
  import jax
  import jax.numpy
  import numpy
  import optax
  from absl.testing import absltest, parameterized
  from jax import random
  from flax import ids
  from flax import linen
  from flax.linen import fp8_ops
  from flax.training import train_state
label_desc: Use image node_id 1 for calling the Fp8Test obj's underlying member method code with example usage: obj.test_fp8_dot_general_injection() and returns: outputs,  grads, jnp