Precision

PrecisionPolicy separates parameter, compute, output, and derivative dtypes. PrecisionMode, PrecisionAlias, and PrecisionName describe accepted precision strings. MatmulPrecision describes the independent JAX dot and convolution precision setting.

Mode	Parameter dtype	Compute dtype	Output dtype	Dynamic loss scaling
64-true	float64	float64	float64	No
32-true	float32	float32	float32	No
16-true	float16	float16	float16	No
bf16-true	bfloat16	bfloat16	bfloat16	No
16-mixed	float32	float16	float32	Yes
bf16-mixed	float32	bfloat16	float32	No

initial_loss_scale applies only to 16-mixed. Non-finite gradients skip the optimizer update and reduce the scale. BFloat16 normally does not need dynamic scaling because it retains a float32-like exponent range.

matmul_precision accepts default, high, or highest. None preserves the current JAX policy. An explicit Trainer value temporarily overrides JAX_DEFAULT_MATMUL_PRECISION during fit and prediction. It does not change parameter or activation dtypes. See JAX's matrix-multiplication precision guide.

PrecisionPolicy [source] dataclass

PrecisionPolicy(
    mode: PrecisionMode,
    parameter_dtype: dtype,
    compute_dtype: dtype,
    output_dtype: dtype,
    derivative_dtype: dtype,
    dynamic_loss_scaling: bool,
    initial_loss_scale: float,
    growth_interval: int,
)

Describe parameter, arithmetic, output, and loss-scaling dtypes.

Mixed modes retain float32 parameters and public model outputs while using a lower precision for dense and embedding arithmetic. True modes use one dtype throughout. This distinction keeps coordinate derivatives and scalar PINN losses in float32 for the recommended mixed-precision modes.

Attributes:

Name	Type	Description
mode	PrecisionMode	Canonical Lightning-compatible precision name.
parameter_dtype	dtype	Data type used to store trainable model parameters.
compute_dtype	dtype	Data type used by supported model arithmetic.
output_dtype	dtype	Data type returned by the model.
derivative_dtype	dtype	Data type used for floating training batches and coordinate derivatives.
dynamic_loss_scaling	bool	Whether finite-gradient dynamic loss scaling is enabled.
initial_loss_scale	float	Initial scale applied before differentiating an FP16 loss.
growth_interval	int	Successful optimizer updates required before increasing the loss scale.

from_name [source] classmethod

from_name(
    precision: PrecisionName | PrecisionPolicy = "32-true",
    *,
    derivative_dtype: Any | None = None,
    initial_loss_scale: float = 32768.0,
    growth_interval: int = 2000,
) -> PrecisionPolicy

Resolve a Lightning-compatible precision mode.

Parameters:

Name	Type	Description	Default
precision	PrecisionName | PrecisionPolicy	Canonical precision mode, supported alias, or existing policy.	'32-true'
derivative_dtype	Any | None	Optional override for floating input batches and coordinate derivatives.	None
initial_loss_scale	float	Initial dynamic scale for 16-mixed training.	32768.0
growth_interval	int	Successful updates between dynamic loss-scale growth events.	2000

Returns:

Type	Description
PrecisionPolicy	Fully resolved immutable precision policy.

Raises:

Type	Description
ValueError	If a mode or loss-scaling option is invalid, or 64-true is requested without JAX x64.

cast_model_state [source]

cast_model_state(model_state: Any) -> Any

Cast floating model-state leaves to the configured parameter dtype.

Parameters:

Name	Type	Description	Default
model_state	Any	Parameter PyTree produced by a supported model.	required

Returns:

Type	Description
Any	Model-state PyTree with floating leaves cast to parameter_dtype.

apply_model_dtype_defaults [source]

apply_model_dtype_defaults(
    model_kwargs: Mapping[str, Any],
) -> dict[str, Any]

Add policy dtypes without replacing explicit architecture options.

Parameters:

Name	Type	Description	Default
model_kwargs	Mapping[str, Any]	Architecture keyword arguments that may explicitly contain parameter_dtype, compute_dtype, or output_dtype.	required

Returns:

Type	Description
dict[str, Any]	New keyword mapping containing policy defaults and all explicit architecture options.

cast_batch [source]

cast_batch(batch: Any) -> Any

Cast floating batch leaves to the configured derivative dtype.

Parameters:

Name	Type	Description	Default
batch	Any	Arbitrary batch PyTree.	required

Returns:

Type	Description
Any	Batch PyTree preserving non-floating leaves and structure.

configure_precision [source]

configure_precision(precision: PrecisionName) -> None

Apply process-wide JAX settings required before array initialization.

Parameters:

Name	Type	Description	Default
precision	PrecisionName	Canonical precision mode or supported alias.	required

Raises:

Type	Description
ValueError	If precision is unknown.