unienv_data.batches.combined_batch¶
CombinedBatch
¶
CombinedBatch(batches: Sequence[BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]], device: Optional[BDeviceType] = None)
Bases: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
get_flattened_at_with_metadata
¶
get_flattened_at_with_metadata(idx: Union[IndexableType, BaseException]) -> Tuple[BArrayType, Dict[str, Any]]
get_at_with_metadata
¶
get_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BatchT, Dict[str, Any]]
set_flattened_at
¶
set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None
append_flattened
¶
append_flattened(value: BArrayType) -> None
Append one flattened sample to 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_single_index_to_batch
¶
convert_single_index_to_batch(backend: ComputeBackend, index_starts: BArrayType, total_length: int, idx: int, device: Optional[BDeviceType] = None) -> Tuple[int, int]
Convert a single index for this batch to a tuple of
- The index of the batch
- The index to index into the batch
convert_index_to_batch
¶
convert_index_to_batch(backend: ComputeBackend, index_starts: BArrayType, total_length: int, idx: Union[IndexableType, BArrayType], device: Optional[BDeviceType] = None) -> Tuple[int, Union[List[Tuple[int, BArrayType, BArrayType]], int]]
Convert an index for this batch to a tuple of
- The length of the resulting array
- List of tuples, each containing:
- The index of the batch
- The index to index into the batch
- The bool mask to index into the resulting array