Skip to content

unienv_data.base.common

IndexableType module-attribute

IndexableType = Union[int, slice, EllipsisType]

BatchT module-attribute

BatchT = TypeVar('BatchT')

SamplerBatchT module-attribute

SamplerBatchT = TypeVar('SamplerBatchT')

SamplerArrayType module-attribute

SamplerArrayType = TypeVar('SamplerArrayType')

SamplerDeviceType module-attribute

SamplerDeviceType = TypeVar('SamplerDeviceType')

SamplerDtypeType module-attribute

SamplerDtypeType = TypeVar('SamplerDtypeType')

SamplerRNGType module-attribute

SamplerRNGType = TypeVar('SamplerRNGType')

BatchBase

BatchBase(single_space: Space[BatchT, BDeviceType, BDtypeType, BRNGType], single_metadata_space: Optional[DictSpace[BDeviceType, BDtypeType, BRNGType]] = None)

Bases: ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]

Base interface for dataset-like collections of structured samples.

BatchBase stores the per-sample space definition and provides a uniform API for retrieving either structured values or their flattened backend-array representation. Concrete subclasses decide where data lives and which mutation operations are supported.

Initialize the batch contract for one logical sample.

is_mutable class-attribute instance-attribute

is_mutable: bool = True

single_space instance-attribute

single_space = single_space

single_metadata_space instance-attribute

single_metadata_space = single_metadata_space

backend property

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

device property

device: Optional[BDeviceType]

get_flattened_at

get_flattened_at(idx: Union[IndexableType, BArrayType]) -> BArrayType

Fetch samples as flattened backend arrays.

get_flattened_at_with_metadata

get_flattened_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BArrayType, Optional[Dict[str, Any]]]

Fetch flattened samples together with optional per-sample metadata.

set_flattened_at

set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None

Overwrite existing samples using flattened data.

append_flattened

append_flattened(value: BArrayType) -> None

Append one flattened sample to the batch.

extend_flattened

extend_flattened(value: BArrayType) -> None

Append a batched block of flattened samples.

get_at

get_at(idx: Union[IndexableType, BArrayType]) -> BatchT

Fetch samples in their structured form.

get_at_with_metadata

get_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BatchT, Optional[Dict[str, Any]]]

Fetch structured samples together with optional metadata.

set_at

set_at(idx: Union[IndexableType, BArrayType], value: BatchT) -> None

Overwrite existing samples using structured data.

remove_at

remove_at(idx: Union[IndexableType, BArrayType]) -> None

Remove one or more samples from the batch.

append

append(value: BatchT) -> None

Append one structured sample to the batch.

extend

extend(value: BatchT) -> None

Append a batched block of structured samples.

extend_from

extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None

Copy data from another batch in bounded-size chunks.

get_slice

get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]

Create a lazy view over a subset of indices.

get_column

get_column(nested_keys: Sequence[str]) -> BatchBase[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]

Create a lazy view over a nested field inside each sample.

close

close() -> None

BatchSampler

BatchSampler(single_space: Space[BatchT, BDeviceType, BDtypeType, BRNGType], single_metadata_space: Optional[DictSpace[BDeviceType, BDtypeType, BRNGType]] = None, batch_size: int = 1)

Bases: Generic[SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType, BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], BatchBase[SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType]

Read-only batch wrapper that samples mini-batches from another batch.

Configure the sampling batch shape without binding the source data yet.

data instance-attribute

data: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]

rng class-attribute instance-attribute

rng: Optional[SamplerRNGType] = None

data_rng class-attribute instance-attribute

data_rng: Optional[BRNGType] = None

is_mutable class-attribute instance-attribute

is_mutable: bool = False

batch_size instance-attribute

batch_size = batch_size

sampled_space property

sampled_space: Space[SamplerBatchT, SamplerDeviceType, SamplerDtypeType, SamplerRNGType]

sampled_metadata_space property

sampled_metadata_space: Optional[DictSpace[SamplerDeviceType, SamplerDtypeType, SamplerRNGType]]

single_space instance-attribute

single_space = single_space

single_metadata_space instance-attribute

single_metadata_space = single_metadata_space

backend property

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

device property

device: Optional[BDeviceType]

manual_seed

manual_seed(seed: int) -> None

Reset sampler RNG state, including the optional data-index RNG.

sample_index

sample_index() -> SamplerArrayType

Draw random indices for one batch.

sample_flat

sample_flat() -> SamplerArrayType

Sample a batch and return it in flattened form.

sample_flat_with_metadata

sample_flat_with_metadata() -> Tuple[SamplerArrayType, Optional[Dict[str, Any]]]

Sample a flattened batch together with optional metadata.

sample

sample() -> SamplerBatchT

Sample a batch in its structured form.

sample_with_metadata

sample_with_metadata() -> Tuple[SamplerBatchT, Optional[Dict[str, Any]]]

Sample a structured batch together with optional metadata.

epoch_iter

epoch_iter() -> Iterator[SamplerBatchT]

Iterate once over a random permutation of the source data.

epoch_iter_with_metadata

epoch_iter_with_metadata() -> Iterator[Tuple[SamplerBatchT, Optional[Dict[str, Any]]]]

Iterate once over shuffled structured batches plus metadata.

epoch_flat_iter

epoch_flat_iter() -> Iterator[SamplerArrayType]

Iterate once over shuffled batches in flattened form.

epoch_flat_iter_with_metadata

epoch_flat_iter_with_metadata() -> Iterator[Tuple[SamplerArrayType, Optional[Dict[str, Any]]]]

Iterate once over shuffled flattened batches plus metadata.

close

close() -> None

get_flattened_at

get_flattened_at(idx: Union[IndexableType, BArrayType]) -> BArrayType

Fetch samples as flattened backend arrays.

get_flattened_at_with_metadata

get_flattened_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BArrayType, Optional[Dict[str, Any]]]

Fetch flattened samples together with optional per-sample metadata.

set_flattened_at

set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None

Overwrite existing samples using flattened data.

append_flattened

append_flattened(value: BArrayType) -> None

Append one flattened sample to the batch.

extend_flattened

extend_flattened(value: BArrayType) -> None

Append a batched block of flattened samples.

get_at

get_at(idx: Union[IndexableType, BArrayType]) -> BatchT

Fetch samples in their structured form.

get_at_with_metadata

get_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BatchT, Optional[Dict[str, Any]]]

Fetch structured samples together with optional metadata.

set_at

set_at(idx: Union[IndexableType, BArrayType], value: BatchT) -> None

Overwrite existing samples using structured data.

remove_at

remove_at(idx: Union[IndexableType, BArrayType]) -> None

Remove one or more samples from the batch.

append

append(value: BatchT) -> None

Append one structured sample to the batch.

extend

extend(value: BatchT) -> None

Append a batched block of structured samples.

extend_from

extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None

Copy data from another batch in bounded-size chunks.

get_slice

get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]

Create a lazy view over a subset of indices.

get_column

get_column(nested_keys: Sequence[str]) -> BatchBase[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]

Create a lazy view over a nested field inside each sample.

convert_index_to_backendarray

convert_index_to_backendarray(backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType], index: IndexableType, length: int, device: Optional[BDeviceType] = None) -> BArrayType

Normalize a Python index into a 1-D backend integer array.

This helper lets batch implementations accept int, slice, and ... while still delegating storage access through a single backend-aware indexing path.