mixed_precision

jaxns.internals.mixed_precision

Module Contents

T

X

class Policy

Encapsulates casting for inputs, outputs and parameters.

measure_dtype: jax.numpy.dtype

index_dtype: jax.numpy.dtype

count_dtype: jax.numpy.dtype

cast_to_index(x, quiet=False)

Converts index values to the index dtype.

Parameters:

• x (X)

• quiet (bool)

Return type:

X

cast_to_measure(x, quiet=False)

Converts measure values to the measure dtype.

Parameters:

• x (X)

• quiet (bool)

Return type:

X

cast_to_count(x, quiet=False)

Converts count values to the count dtype.

Parameters:

• x (X)

• quiet (bool)

Return type:

X

mp_policy

float_type

int_type

complex_type