unienv_data.base.common¶
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.
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.
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.
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.
sampled_space
property
¶
sampled_space: Space[SamplerBatchT, SamplerDeviceType, SamplerDtypeType, SamplerRNGType]
sampled_metadata_space
property
¶
sampled_metadata_space: Optional[DictSpace[SamplerDeviceType, SamplerDtypeType, SamplerRNGType]]
manual_seed
¶
manual_seed(seed: int) -> None
Reset sampler RNG state, including the optional data-index RNG.
sample_flat_with_metadata
¶
sample_flat_with_metadata() -> Tuple[SamplerArrayType, Optional[Dict[str, Any]]]
Sample a flattened batch together with optional metadata.
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.
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.
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.