jmp

import jax import jax.numpy as jnp import jmp # Define floating point types based on your hardware (e.g., bfloat16 for TPU, float16 for GPU) half = jnp.float16 # or jnp.bfloat16 for TPUs full = jnp.float32 # Create a mixed precision policy # Parameters stored in full precision, computation and output in half precision policy = jmp.Policy(param_dtype=full, compute_dtype=half, output_dtype=half) # Example: Applying policy to a JAX array x = jnp.array([1.0, 2.0, 3.0], dtype=full) x_half = policy.cast_to_compute(x) print(f"Original (full): {x.dtype}, {x}") print(f"Computed (half): {x_half.dtype}, {x_half}") # Example: Using DynamicLossScale # Initialize DynamicLossScale with an initial loss scale value loss_scale = jmp.DynamicLossScale(jmp.half_dtype()(2**15)) # Using a jmp dtype helper # Simulate a gradient check and adjustment grads_finite = True # Assume gradients were finite in this step loss_scale = loss_scale.adjust(grads_finite) print(f"Adjusted loss scale: {loss_scale.loss_scale}")