Skip to content

abstract_dataloader.ext.types

Data type system framework following the ADL recommendations.

Programming Model

  • Declare types as dataclasses using the provided dataclass decorator, which registers them to the global optree namespace.
  • Set each class as Generic[TArray], where TArray is a TypeVar which is ArrayLike, e.g., torch.Tensor, jax.Array or np.ndarray.

Warning

This module requires optree to be installed.

abstract_dataloader.ext.types.TArray module-attribute

TArray = TypeVar('TArray', bound=ArrayLike)

Type variable for ArrayLike types.

abstract_dataloader.ext.types.ArrayLike

Bases: Protocol

Array type, e.g., torch.Tensor | jax.Array | np.ndarray.

Use this type to specify arbitrary array types.

Source code in src/abstract_dataloader/ext/types.py
@runtime_checkable
class ArrayLike(Protocol):
    """Array type, e.g., `torch.Tensor | jax.Array | np.ndarray`.

    Use this type to specify arbitrary array types.
    """

    @property
    def shape(self) -> tuple[int, ...]: ...

    @property
    def dtype(self) -> Any: ...

abstract_dataloader.ext.types.dataclass

dataclass(cls)

A dataclass decorator which registers into optree's global namespace.

Source code in src/abstract_dataloader/ext/types.py
@dataclass_transform(field_specifiers=(field,))
def dataclass(cls):  # noqa: D103
    """A dataclass decorator which registers into optree's global namespace."""
    return _optree_dataclass(cls, namespace='')