code: def setup(self) -> None:
      scale_args = (
          initializers.ones_init(),
          random.PRNGKey(0),
          (1,),
          jnp.float32,
      )
      amax_history_args = (
          initializers.zeros_init(),
          random.PRNGKey(0),
          (self.amax_history_length,),
          jnp.float32,
      )
  
      self.input_amax_history = self.variable(
          OVERWRITE_WITH_GRADIENT,
          "input_amax_history",
          *amax_history_args,
      )
      self.kernel_amax_history = self.variable(
          OVERWRITE_WITH_GRADIENT,
          "kernel_amax_history",
          *amax_history_args,
      )
      self.output_grad_amax_history = self.variable(
          OVERWRITE_WITH_GRADIENT,
          "output_grad_amax_history",
          *amax_history_args,
      )
  
      self.input_scale = self.variable(
          OVERWRITE_WITH_GRADIENT, "input_scale", *scale_args
      )
      self.kernel_scale = self.variable(
          OVERWRITE_WITH_GRADIENT, "kernel_scale", *scale_args
      )
      self.output_grad_scale = self.variable(
          OVERWRITE_WITH_GRADIENT, "output_grad_scale", *scale_args
      )
  
imports: import warnings
  from functools import partial
  from jax import custom_jvp, custom_vjp, lax, random
  from jax import numpy
  from flax.linen import initializers, module
label_desc: Use image node_id 1 for calling the Fp8DotGeneralOp obj's underlying member method code with example usage: obj.setup() without return types