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 (jax.typing.ArrayLike) – an array or scalar value to be cast
new_dtype (jax.typing.DTypeLike) – a NumPy dtype representing the target type.
- Returns:
An array with the same shape as operand, cast elementwise to new_dtype.
- Return type: