kd.data.Cast

kd.data.Cast#

class kauldron.data.Cast(*, key: typing.Annotated[typing.Any, <object object at 0x76412092fb90>] | typing.Sequence[typing.Annotated[typing.Any, <object object at 0x76412092fb90>]] | dict[typing.Annotated[typing.Any, <object object at 0x76412092fb90>], typing.Annotated[typing.Any, <object object at 0x76412092fb90>]], dtype: typing.Any)[source]

Bases: kauldron.data.transforms.base.ElementWiseTransform

Cast an element to the specified dtype.

dtype: Any
map_element(
element: jaxtyping.Shaped[Tensor, '*any'] | jaxtyping.Shaped[ndarray, '*any'] | jaxtyping.Shaped[Array, '*any'],
) → jaxtyping.Shaped[Tensor, '*any'] | jaxtyping.Shaped[ndarray, '*any'] | jaxtyping.Shaped[Array, '*any'][source]