jax.lax.convert_element_type

Contents

jax.lax.convert_element_type#

jax.lax.convert_element_type(operand, new_dtype)[source]#

Elementwise cast.

Wraps XLA’s ConvertElementType operator, which performs an elementwise conversion from one type to another. Similar to a C++ static_cast.

Parameters:

operand (Union[Array, ndarray, bool_, number, bool, int, float, complex]) – an array or scalar value to be cast
new_dtype (Union[str, type[Any], dtype, SupportsDType]) – a NumPy dtype representing the target type.

Return type:

Returns:

An array with the same shape as operand, cast elementwise to new_dtype.