code: def quantize(x, q_dtype, scale, compute_dtype):
      # Explicitly cast the max values to the compute dtype to avoid unnecessary
      # casting to FP32 during the subsequent math operations."
      dtype_max = get_fp8_max(q_dtype, compute_dtype)
      scaled_x = x / jnp.broadcast_to(
          scale.astype(compute_dtype), x.shape
      )
      clipped_x = jnp.clip(scaled_x, -dtype_max, dtype_max)
      return clipped_x.astype(q_dtype)
  
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 2 for calling a global function with example usage: quantize(x, q_dtype, scale, compute_dtype) and returns: clipped_x