kd.data.Rearrange#

class kauldron.data.Rearrange(*, 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>]], pattern: str, axes_lengths: dict[str, int] = <factory>)[source]

Bases: kauldron.data.transforms.base.ElementWiseTransform

Einops rearrange on a single element.

Mostly a wrapper around einops.rearrange, but also supports basic types like int, float, lists and tuples (which are converted to a numpy array first).

Example:

cfg.train_ds = kd.data.tf.Tfds(
    ...
    transforms=[
        ...,
        kd.data.Rearrange(key="image", pattern="h w c -> (h w c)"),
    ]
)

pattern

einops.rearrange pattern, e.g. “b h w c -> b c (h w)”

Type:: str

axes_lengths

a dictionary for specifying additional axis sizes that cannot be inferred from the pattern and the tensor alone.

Type:: dict[str, int]

pattern: str

axes_lengths: dict[str, int]

map_element( element: Any, ) → kauldron.typing.array_types.XArray[source]