kd.data.py.Rearrange

class kauldron.data.py.Rearrange(*, key: typing.Annotated[typing.Any, <object object at 0x7824c478ba80>] | typing.Sequence[typing.Annotated[typing.Any, <object object at 0x7824c478ba80>]] | dict[typing.Annotated[typing.Any, <object object at 0x7824c478ba80>], typing.Annotated[typing.Any, <object object at 0x7824c478ba80>]], pattern: str, axes_lengths: dict[str, int] = <factory>)

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

key: kontext.Key | Sequence[kontext.Key] | dict[kontext.Key, kontext.Key]