code: def get_fp8_max(fp8_dtype, out_dtype):
      assert fp8_dtype in (jnp.float8_e4m3fn, jnp.float8_e5m2)
      return jnp.finfo(fp8_dtype).max.astype(out_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 1 for calling a global function with example usage: get_fp8_max(fp8_dtype, out_dtype) and returns: jnp