Skip to content

TensorStore

API reference

API reference¶

Core¶

class tensorstore.TensorStore: Asynchronous multi-dimensional array handle.

class tensorstore.Context: Manages shared TensorStore context resources, such as caches and credentials.

class tensorstore.KvStore: Key-value store that maps an ordered set of byte string keys to byte string values.

class tensorstore.Transaction: Transactions are used to stage a group of modifications (e.g. writes to tensorstore.TensorStore objects) in memory, and then either commit the group all at once or abort it.

class tensorstore.Batch: Batches are used to group together read operations for potentially improved efficiency.

tensorstore.open(spec: Spec | Any, *, ...) → Future[TensorStore]: Opens or creates a TensorStore from a Spec.

Indexing¶

tensorstore.inf : int = 4611686018427387903: Special constant equal to \(2^{62}-1\) that indicates an unbounded index domain.

class tensorstore.IndexDomain: Domain (including bounds and optional dimension labels) of an N-dimensional index space.

class tensorstore.IndexTransform: Represents a transform from an input index space to an output space.

class tensorstore.Dim: 1-d index interval with optionally-implicit bounds and dimension label.

class tensorstore.OutputIndexMap: Represents an output index map for an index transform.

class tensorstore.OutputIndexMaps: View of the output index maps for an index transform.

class tensorstore.OutputIndexMethod: Indicates the output index method of an OutputIndexMap.

class tensorstore.DimExpression: Specifies an advanced indexing operation.

class tensorstore.d(DimExpression): Specifies a dimension selection, for starting a dimension expression.

tensorstore.newaxis = None: Alias for None used in indexing expressions to specify a new singleton dimension.

class tensorstore.Indexable: Abstract base class for types that support TensorStore indexing operations.

Spec¶

class tensorstore.Spec: Specification for opening or creating a TensorStore.

class tensorstore.Schema: Driver-independent options for defining a TensorStore schema.

class tensorstore.CodecSpec: Specifies driver-specific encoding/decoding parameters.

class tensorstore.OpenMode: Specifies the mode to use when opening a TensorStore.

class tensorstore.ChunkLayout: Describes the storage layout of a tensorstore.TensorStore.

class tensorstore.Unit: Specifies a physical quantity/unit.

tensorstore.RecheckCacheOption: Determines under what circumstances cached data is revalidated.

Views¶

tensorstore.cast(base: TensorStore, dtype: dtype) → TensorStore: Returns a read/write view with the data type converted.

tensorstore.cast(base: Spec, dtype: dtype) → Spec: Returns a view with the data type converted.

tensorstore.array(array: ArrayLike, ...) → TensorStore: Returns a TensorStore that reads/writes from an in-memory array.

tensorstore.overlay(layers, ...) → TensorStore: Virtually overlays a sequence of TensorStore layers within a common domain.

tensorstore.stack(layers, ...) → TensorStore: Virtually stacks a sequence of TensorStore layers along a new dimension.

tensorstore.concat(layers, ...) → TensorStore: Virtually concatenates a sequence of TensorStore layers along an existing dimension.

tensorstore.downsample(base: TensorStore, ...) → TensorStore: Returns a virtual downsampled view of a TensorStore.

tensorstore.downsample(base: Spec, downsample_factors, ...) → Spec: Returns a virtual downsampled view view of a Spec.

Virtual views¶

class tensorstore.VirtualChunkedReadParameters: Options passed to read callbacks used with virtual_chunked.

class tensorstore.VirtualChunkedWriteParameters: Options passed to write callbacks used with virtual_chunked.

tensorstore.virtual_chunked(...) → TensorStore: Creates a TensorStore where the content is read/written chunk-wise by an arbitrary function.

Data types¶

class tensorstore.dtype: TensorStore data type representation.

tensorstore.bool : dtype = dtype("bool"): Boolean data type (0 or 1). Corresponds to the bool type and numpy.bool_.

tensorstore.int4 : dtype = dtype("int4"): 4-bit signed two’s-complement integer data type, internally stored as its 8-bit signed integer equivalent (i.e. sign-extended). Corresponds to jax.numpy.int4.

tensorstore.int8 : dtype = dtype("int8"): 8-bit signed two’s-complement integer data type. Corresponds to numpy.int8.

tensorstore.uint8 : dtype = dtype("uint8"): 8-bit unsigned integer. Corresponds to numpy.uint8.

tensorstore.int16 : dtype = dtype("int16"): 16-bit signed two’s-complement integer data type. Corresponds to numpy.int16.

tensorstore.uint16 : dtype = dtype("uint16"): 16-bit unsigned integer. Corresponds to numpy.uint16.

tensorstore.int32 : dtype = dtype("int32"): 32-bit signed two’s-complement integer data type. Corresponds to numpy.int32.

tensorstore.uint32 : dtype = dtype("uint32"): 32-bit unsigned integer. Corresponds to numpy.uint32.

tensorstore.int64 : dtype = dtype("int64"): 32-bit signed two’s-complement integer data type. Corresponds to numpy.int64.

tensorstore.uint64 : dtype = dtype("uint64"): 64-bit unsigned integer data type. Corresponds to numpy.uint64.

tensorstore.float8_e4m3fn : dtype = dtype("float8_e4m3fn"): 8-bit floating-point data type.

tensorstore.float8_e4m3fnuz : dtype = dtype("float8_e4m3fnuz"): 8-bit floating-point data type.

tensorstore.float8_e4m3b11fnuz : dtype = dtype("float8_e4m3b11fnuz"): 8-bit floating-point data type.

tensorstore.float8_e5m2 : dtype = dtype("float8_e5m2"): 8-bit floating-point data type.

tensorstore.float8_e5m2fnuz : dtype = dtype("float8_e5m2fnuz"): 8-bit floating-point data type.

tensorstore.float16 : dtype = dtype("float16"): IEEE 754 binary16 half-precision floating-point data type. Correspond to numpy.float16.

tensorstore.bfloat16 : dtype = dtype("bfloat16"): bfloat16 floating-point data type.

tensorstore.float32 : dtype = dtype("float32"): IEEE 754 binary32 single-precision floating-point data type. Corresponds to numpy.float32.

tensorstore.float64 : dtype = dtype("float64"): IEEE 754 binary64 double-precision floating-point data type. Corresponds to numpy.float64.

tensorstore.complex64 : dtype = dtype("complex64"): Complex number based on float32. Corresponds to numpy.complex64.

tensorstore.complex128 : dtype = dtype("complex128"): Complex number based on float64. Corresponds to numpy.complex128.

tensorstore.string : dtype = dtype("string"): Variable-length byte string data type. Corresponds to the Python bytes type.

tensorstore.ustring : dtype = dtype("ustring"): Variable-length Unicode string data type. Corresponds to the Python str type.

tensorstore.json : dtype = dtype("json"): JSON data type. Corresponds to an arbitrary Python JSON value.

Asynchronous support¶

class tensorstore.Future: Handle for consuming the result of an asynchronous operation.

class tensorstore.Promise: Handle for producing the result of an asynchronous operation.

class tensorstore.WriteFutures: Handle for consuming the result of an asynchronous write operation.

class tensorstore.FutureLike: Abstract base class for types representing an asynchronous result.

OCDBT¶

class tensorstore.ocdbt.DistributedCoordinatorServer: Distributed coordinator server for the OCDBT (Optionally-Cooperative Distributed B+Tree) database.

tensorstore.ocdbt.dump(base: KvStore, ...) → Future[Any]: Dumps the internal representation of an OCDBT database.

Experimental¶

tensorstore.experimental_collect_matching_metrics(...) → list[Any]: Collects metrics with a matching prefix.

tensorstore.experimental_collect_prometheus_format_metrics(...) → list[str]: Collects metrics in prometheus exposition format. See: https://prometheus.io/docs/instrumenting/exposition_formats/

tensorstore.experimental_push_metrics_to_prometheus(...) → Future[int]: Publishes metrics to the prometheus pushgateway. See: https://github.com/prometheus/pushgateway

tensorstore.experimental_update_verbose_logging(...) → None: Updates verbose logging flags associated with –tensorstore_verbose_logging and TENSORSTORE_VERBOSE_LOGGING flags.

tensorstore.parse_tensorstore_flags(argv: list[str]) → None: Parses and initializes internal tensorstore flags from argv.