code: def compute_scale_and_amax_history(x, q_dtype, scale, amax_history):
      dtype_max = get_fp8_max(q_dtype, jnp.float32)
      amax_update = jnp.max(jnp.abs(x)).astype(scale.dtype)
      new_history = (
          jnp.roll(amax_history, shift=-1, axis=0)
          .at[0]
          .set(amax_update)
      )
      amax_from_history = jnp.max(new_history, axis=0)
      new_scale = compute_scale(amax_from_history, scale, dtype_max)
      return new_scale, new_history
  
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 6 for calling a global function with example usage: compute_scale_and_amax_history(x, q_dtype, scale, amax_history) and returns: new_scale,  new_history