LazyStackedTensorDictStore#

LazyStackedTensorDictStore has the methods of TensorDictBase, listed by task in Methods by task. This page lists only the methods that LazyStackedTensorDictStore adds, and the overrides that have their own documentation.

class tensordict.store.LazyStackedTensorDictStore(*, backend: Literal['redis', 'dragonfly'] = 'redis', host: str = 'localhost', port: int = 6379, db: int = 0, unix_socket_path: str | None = None, prefix: str = 'tensordict', count: int, stack_dim: int = 0, inner_batch_size: Sequence[int], device=None, client=None, td_id: str | None = None, cache_metadata: bool = True, **redis_kwargs)#

Bases: _StoreConnection, _StoreUnsupported, TensorDictBase

A lazy stack of tensordicts whose entries live on a Redis-compatible server.

Warning

LazyStackedTensorDictStore is experimental: its API can change in a minor release without a deprecation cycle.

TensorDictStore.from_tensordict returns this class when you pass a LazyStackedTensorDict, so that the elements of the stack can have entries of different shapes. You can also call from_lazy_stack(), and from_store() to reconnect. The class behaves like a lazy stack, but it is a TensorDictBase subclass, not a LazyStackedTensorDict subclass.

Supports Redis, Dragonfly, KeyDB, and any other Redis-wire-compatible server. The data lifetime, close and clear_redis work as for TensorDictStore (see its “Lifecycle and limits” section).

Warning

store["nested"] is an in-memory copy of the nested entry: writes through it, for example store["nested"].set_(...) or store["nested"].zero_(), change the copy only, not the store. Write with the full key, store.set(("nested", "b"), ...), or with update().

store[i] is a view of element i that reads and writes on the server. On a locked store, the view follows the lock of the store: it writes the keys that the store has, with or without inplace, and a write of a new key raises the lock error. Unlike stack[i] of a locked lazy stack, the view itself is not locked: its is_locked is False, and the results of operations on it, such as store[i] + 1 or store[i].to("cpu"), are unlocked tensordicts.

Each leaf key is one string on the server for the whole stack, with a table of offsets when the shapes of the elements differ, so the number of keys on the server does not grow with the number of elements.

The constructor makes an empty stack of count elements of batch size inner_batch_size. To copy a lazy stack to the server, use from_lazy_stack().

Keyword Arguments:

backend (str, optional) – a label for the kind of server, as in TensorDictStore. Defaults to "redis".
host (str, optional) – the host name of the server. Defaults to "localhost".
port (int, optional) – the port of the server. Defaults to 6379.
db (int, optional) – the database number on the server. Defaults to 0.
unix_socket_path (str, optional) – the path of a Unix socket, for a server on the same machine. If set, host and port are ignored. Defaults to None.
prefix (str, optional) – the namespace of the keys of the store on the server. Defaults to "tensordict".
count (int) – the number of elements of the stack.
stack_dim (int, optional) – the dimension of the stack in the batch size. Defaults to 0.
inner_batch_size (sequence of int) – the batch size of each element.
device (device, str or int, optional) – the device of the tensors that reads return. Defaults to None: reads return CPU tensors, and the store has no device.
client (redis.asyncio.Redis, optional) – an existing client. If given, host, port, db and unix_socket_path are ignored. Defaults to None, which creates a client.
td_id (str, optional) – the identifier of the store on the server. Defaults to None, which creates a new UUID. To reconnect to an existing store, use from_store().
cache_metadata (bool, optional) – if True, keeps the shape and dtype of each entry in the process, to save round-trips to the server. Defaults to True.
**redis_kwargs – passed to redis.asyncio.Redis.

Examples

The example needs a Redis server on localhost:6379.

>>> import torch
>>> from tensordict import TensorDict, lazy_stack
>>> from tensordict.store import TensorDictStore
>>> local = lazy_stack([TensorDict(x=torch.zeros(3)), TensorDict(x=torch.zeros(5))])
>>> store = TensorDictStore.from_tensordict(local)  
>>> type(store).__name__  
'LazyStackedTensorDictStore'
>>> store[0]["x"].shape, store[1]["x"].shape  
(torch.Size([3]), torch.Size([5]))
>>> store[1] = TensorDict(x=torch.ones(5))  
>>> store[1]["x"]  
tensor([1., 1., 1., 1., 1.])
>>> store.is_locked = True  
>>> store.is_locked, store[1].is_locked  
(True, False)
>>> store[1]["x"] = torch.zeros(5)  
>>> store[1]["y"] = torch.zeros(5)  
Traceback (most recent call last):
...
RuntimeError: Cannot modify locked TensorDict. ...
>>> store.is_locked = False  
>>> store.clear_redis()  
>>> store.close()  

flatten_keys(separator='.', inplace=False, is_leaf=None)#

Returns a local copy of the store whose nested keys are joined by separator.

The data is read from the server with to_tensordict(), and the store is not changed.

Parameters:

separator (str, optional) – the string that joins the parts of a nested key. Defaults to ".".
inplace (bool, optional) – must be False; True raises a ValueError. Defaults to False.
is_leaf (Callable, optional) – see flatten_keys(). Defaults to None.

See flatten_keys().

classmethod from_dict(input_dict, *, auto_batch_size: bool = False, batch_size=None, device=None, **kwargs)#

Raises a NotImplementedError: build the store with from_lazy_stack().

The arguments are those of from_dict(), and are not used.

classmethod from_lazy_stack(lazy_td, *, backend: Literal['redis', 'dragonfly'] = 'redis', host: str = 'localhost', port: int = 6379, db: int = 0, unix_socket_path: str | None = None, prefix: str = 'tensordict', device=None, **kwargs) → LazyStackedTensorDictStore#

Returns a new store that holds a copy of the data of a lazy stack.

The data is sent in chunks of 10,000 elements, so that the full stack is not built in memory at once.

Parameters:

lazy_td (LazyStackedTensorDict) – the lazy stack to copy to the server.

Keyword Arguments:

backend (str, optional) – a label for the kind of server, as in TensorDictStore. Defaults to "redis".
host (str, optional) – the host name of the server. Defaults to "localhost".
port (int, optional) – the port of the server. Defaults to 6379.
db (int, optional) – the database number on the server. Defaults to 0.
unix_socket_path (str, optional) – the path of a Unix socket, for a server on the same machine. If set, host and port are ignored. Defaults to None.
prefix (str, optional) – the namespace of the keys of the store on the server. Defaults to "tensordict".
device (device, str or int, optional) – the device of the tensors that reads return. Defaults to None, which takes the device of lazy_td.
**kwargs – passed to the constructor, such as client or td_id, and options of the Redis client.

Raises:

TypeError – if lazy_td is not a LazyStackedTensorDict.

Examples

The example needs a Redis server on localhost:6379.

>>> import torch
>>> from tensordict import TensorDict, lazy_stack
>>> from tensordict.store import LazyStackedTensorDictStore
>>> local = lazy_stack([TensorDict(x=torch.zeros(2)), TensorDict(x=torch.zeros(4))])
>>> store = LazyStackedTensorDictStore.from_lazy_stack(local)  
>>> store.batch_size, store[1]["x"].shape  
(torch.Size([2]), torch.Size([4]))
>>> store.clear_redis()  
>>> store.close()  

classmethod from_store(*, backend: Literal['redis', 'dragonfly'] = 'redis', host: str = 'localhost', port: int = 6379, db: int = 0, unix_socket_path: str | None = None, prefix: str = 'tensordict', td_id: str, device=None, **kwargs) → LazyStackedTensorDictStore#

Returns a store connected to a lazy stack that another store wrote to the server.

Reads the number of elements, the stack dimension, the batch size of the elements and the device from the server.

Keyword Arguments:

backend (str, optional) – a label for the kind of server, as in TensorDictStore. Defaults to "redis".
host (str, optional) – the host name of the server. Defaults to "localhost".
port (int, optional) – the port of the server. Defaults to 6379.
db (int, optional) – the database number on the server. Defaults to 0.
unix_socket_path (str, optional) – the path of a Unix socket, for a server on the same machine. If set, host and port are ignored. Defaults to None.
prefix (str, optional) – the namespace of the keys of the store on the server. Defaults to "tensordict".
td_id (str) – the identifier of the store to connect to. A store shows it in its repr, and holds it in the private attribute _td_id; there is no public accessor.
device (device, str or int, optional) – the device of the tensors that reads return. Defaults to None, which uses the device stored on the server.
**kwargs – passed to the constructor, such as client or cache_metadata, and options of the Redis client.

Returns:

A LazyStackedTensorDictStore connected to the existing data.

Raises:

KeyError – if the server has no lazy stack with this td_id.

Examples

The example needs a Redis server on localhost:6379. The second store could be in another process.

>>> import torch
>>> from tensordict import TensorDict, lazy_stack
>>> from tensordict.store import LazyStackedTensorDictStore
>>> local = lazy_stack([TensorDict(x=torch.zeros(2)), TensorDict(x=torch.zeros(4))])
>>> writer = LazyStackedTensorDictStore.from_lazy_stack(local)  
>>> reader = LazyStackedTensorDictStore.from_store(td_id=writer._td_id)  
>>> reader[1]["x"].shape  
torch.Size([4])
>>> writer.clear_redis()  
>>> writer.close()  
>>> reader.close()  

to_local() → TensorDict#: Returns a local copy of the data of the store.

Deprecated since version 0.15: Use to_tensordict() instead.

LazyStackedTensorDictStore#

Docs

Tutorials

Resources