TensorDict

https://docs.pytorch.org/tensordict/stable/
Admin
TensorDict is a dictionary-like class that inherits properties from tensors, enabling readable and...
Tokens:613,456
Snippets:3,403
Trust Score:10
Update:9 months ago
Show doc for...
Tokens:
Raw
### Full Example: Saving and Loading Memory-Mapped TensorDict

Source: https://docs.pytorch.org/tensordict/stable/saving

Demonstrates a complete workflow for saving and loading a memory-mapped TensorDict using torchsnapshot. This example includes creating, memmapping, saving, and restoring a TensorDict, ensuring data integrity and type preservation.

```python
>>> td = TensorDict({"a": torch.randn(3), "b": TensorDict({"c": torch.randn(3, 1)}, [3, 1])}, [3])
>>> td.memmap_()
>>> assert isinstance(td["b", "c"], MemmapTensor)
>>>
>>> app_state = {
...     "state": torchsnapshot.StateDict(tensordict=td.state_dict(keep_vars=True))
... }
>>> snapshot = torchsnapshot.Snapshot.take(app_state=app_state, path=f"/tmp/{uuid.uuid4()}")
>>>
>>>
>>> td_dest = TensorDict({"a": torch.zeros(3), "b": TensorDict({"c": torch.zeros(3, 1)}, [3, 1])}, [3])
>>> td_dest.memmap_()
>>> assert isinstance(td_dest["b", "c"], MemmapTensor)
>>> app_state = {
...     "state": torchsnapshot.StateDict(tensordict=td_dest.state_dict(keep_vars=True))
... }
>>> snapshot.restore(app_state=app_state)
>>> # sanity check
>>> assert (td_dest == td).all()
>>> assert (td_dest["b"].batch_size == td["b"].batch_size)
>>> assert isinstance(td_dest["b", "c"], MemmapTensor)

```

--------------------------------

### Example: TensorDict Keys Usage

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.LazyStackedTensorDict

Demonstrates how to use the keys() method with different options to retrieve keys from a TensorDict. Shows examples for default behavior, leaves_only, and include_nested with leaves_only.

```python
from tensordict import TensorDict
data = TensorDict({"0": 0, "1": {"2": 2}}, batch_size=[])
print(data.keys())
print(list(data.keys(leaves_only=True)))
print(list(data.keys(include_nested=True, leaves_only=True)))
```

--------------------------------

### Example TensorDict Initialization

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.PersistentTensorDict

Example demonstrating the initialization of a `tensorclass` and its usage with PyTorch tensors and integers. This showcases a basic structure for creating custom tensor-compatible classes.

```python
>>> from tensordict import TensorDict, tensorclass
>>> import torch
>>>
>>> @tensorclass
>>> class MyClass:
...     x: torch.Tensor
...     y: int
>>>
```

--------------------------------

### Example: CompositeDistribution Usage with TensorDict

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.nn.distributions.CompositeDistribution

Demonstrates how to create and use a CompositeDistribution with a TensorDict. This example shows parameter setup, distribution mapping, sampling, and calculating log probabilities with aggregation disabled.

```python
>>> params = TensorDict({
...     "cont": {"loc": torch.randn(3, 4), "scale": torch.rand(3, 4)},
...     ("nested", "disc"): {"logits": torch.randn(3, 10)}
... }, [3])
>>> dist = CompositeDistribution(params,
...     distribution_map={"cont": d.Normal, ("nested", "disc"): d.Categorical})
>>> sample = dist.sample((4,))
>>> with set_composite_lp_aggregate(False):
...     sample = dist.log_prob(sample)
...     print(sample)
TensorDict(
    fields={
        cont: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        cont_log_prob: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        nested: TensorDict(
            fields={
                disc: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.int64, is_shared=False),
                disc_log_prob: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)

```

--------------------------------

### Example: Loading a Memory-Mapped TensorDict

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.LazyStackedTensorDict

Demonstrates loading a TensorDict from disk using `load_memmap`. It also shows how to assert the equality of the loaded TensorDict with the original.

```python
from tensordict import TensorDict
td = TensorDict.fromkeys(["a", "b", "c", ("nested", "e")])
td.memmap("./saved_td")
td_load = TensorDict.load_memmap("./saved_td")
assert (td == td_load).all()
```

--------------------------------

### LazyStackedTensorDict Examples

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.LazyStackedTensorDict

Illustrates how to create and use LazyStackedTensorDict, including stacking TensorDicts, accessing elements, and performing assignments.

```APIDOC
## LazyStackedTensorDict Examples

### Basic Stacking and Access
```python
from tensordict import TensorDict
import torch
tds = [TensorDict({'a': torch.randn(3, 4)}, batch_size=[3])
    for _ in range(10)]
td_stack = torch.stack(tds, -1)
print(td_stack.shape)
# Output: torch.Size([3, 10])
print(td_stack.get("a").shape)
# Output: torch.Size([3, 10, 4])
print(td_stack[:, 0] is tds[0])
# Output: True
```

### Assignment with Lists
```python
td = LazyStackedTensorDict(TensorDict(), TensorDict(), stack_dim=0)
td["a"] = [torch.ones(2), torch.zeros(1)]
assert td[1]["a"] == torch.zeros(1)
td["b"] = ["a string", "another string"]
assert td[1]["b"] == "another string"
```

### Controlling `get()` Method Behavior
When using the `get()` method, one can pass `as_nested_tensor`, `as_padded_tensor` or the `as_list` arguments to control how the data should be presented if the dimensions of the tensors mismatch. When passed, the nesting/padding will occur regardless of whether the dimensions mismatch or not.
```

--------------------------------

### Python Multiprocessing Queue Example

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.NonTensorStack

Demonstrates using a multiprocessing Queue to pass data between a server and a client process. The client waits for a signal from the server, and the example includes error handling and process joining with timeouts.

```python
import multiprocessing as mp

def server(queue):
    # Simulate some work
    queue.put("yuppie")

def client(queue):
    # Client process would typically do something else
    pass

if __name__ == "__main__":
    queue = mp.Queue(1)

    main_worker = mp.Process(target=server, args=(queue,))
    secondary_worker = mp.Process(target=client, args=(queue,))

    main_worker.start()
    secondary_worker.start()

    try:
        out = queue.get(timeout=10)
        print(out)
    finally:
        queue.close()
        main_worker.join(timeout=10)
        secondary_worker.join(timeout=10)
```

--------------------------------

### Example: Loading Nested TensorDicts

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.LazyStackedTensorDict

Illustrates loading a nested TensorDict from a specific subdirectory within the saved path.

```python
nested = TensorDict.load_memmap("./saved_td/nested")
assert nested["e"] == 0
```

--------------------------------

### Full Example: Saving and Restoring TensorDict with Memmap

Source: https://docs.pytorch.org/tensordict/stable/_sources/saving.rst

A comprehensive example demonstrating the saving and restoring of a TensorDict with memmap tensors using torchsnapshot. This includes setup, saving, restoration, and verification steps.

```python
>>> td = TensorDict({"a": torch.randn(3), "b": TensorDict({"c": torch.randn(3, 1)}, [3, 1])}, [3])
>>> td.memmap_()
>>> assert isinstance(td["b", "c"], MemmapTensor)
>>> 
>>> app_state = {
...     "state": torchsnapshot.StateDict(tensordict=td.state_dict(keep_vars=True))
... }
>>> snapshot = torchsnapshot.Snapshot.take(app_state=app_state, path=f"/tmp/{uuid.uuid4()}")
>>> 
>>> 
>>> td_dest = TensorDict({"a": torch.zeros(3), "b": TensorDict({"c": torch.zeros(3, 1)}, [3, 1])}, [3])
>>> td_dest.memmap_()
>>> assert isinstance(td_dest["b", "c"], MemmapTensor)
>>> app_state = {
...     "state": torchsnapshot.StateDict(tensordict=td_dest.state_dict(keep_vars=True))
... }
>>> snapshot.restore(app_state=app_state)
>>> # sanity check
>>> assert (td_dest == td).all()
>>> assert (td_dest["b"].batch_size == td["b"].batch_size)
>>> assert isinstance(td_dest["b", "c"], MemmapTensor)
```

--------------------------------

### Preallocate TensorDict Dataset on Disk - Python

Source: https://docs.pytorch.org/tensordict/stable/saving

Demonstrates how to preallocate a large dataset structure on disk using `memmap_like`. This is useful for datasets where individual data points are created on the fly, and the overall structure needs to be defined on disk without immediate memory allocation for all data.

```python
>>> import torch
>>> from tensordict import TensorDict
>>> def make_datum(): # used for illustration purposes
...    return TensorDict({"image": torch.randint(255, (3, 64, 64)), "label": 0}, batch_size=[])
>>> dataset_size = 1_000_000
>>> datum = make_datum() # creates a single instance of a TensorDict datapoint
>>> data = datum.expand(dataset_size) # does NOT require more memory usage than datum, since it's only a view on datum!
>>> data_disk = data.memmap_like("/path/to/data")  # creates the two memory-mapped tensors on disk
>>> del data # data is not needed anymore
```

--------------------------------

### Reshape TensorDict Example

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.LazyStackedTensorDict

Demonstrates how to reshape a TensorDict and its contained tensors. This operation modifies the shape of the tensor associated with the 'x' key to a flattened 1D tensor. Ensure PyTorch is installed for tensor operations.

```python
>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict({
...     'x': torch.arange(12).reshape(3, 4),
... }, batch_size=[3, 4])
>>> td = td.reshape(12)
>>> print(td['x'])
...
```

--------------------------------

### TensorDictModule Example with Probabilistic Distribution

Source: https://docs.pytorch.org/tensordict/stable/_downloads/ba42555713e16dc342a5f3e550dec192/tensordict_module

Demonstrates the use of ProbabilisticTensorDictModule and TensorDictSequential to create a neural network pipeline that processes inputs, extracts distribution parameters, and samples from a Normal distribution. It shows how the TensorDict is modified in-place.

```python
from tensordict.nn import (
    ProbabilisticTensorDictModule,
    ProbabilisticTensorDictSequential,
)
from tensordict.nn.distributions import NormalParamExtractor
from torch import distributions as dist

td = TensorDict({"input": torch.randn(3, 4), "hidden": torch.randn(3, 8)}, [3])
net = torch.nn.GRUCell(4, 8)
net = TensorDictModule(net, in_keys=["input", "hidden"], out_keys=["hidden"])
extractor = NormalParamExtractor()
extractor = TensorDictModule(extractor, in_keys=["hidden"], out_keys=["loc", "scale"])
td_module = ProbabilisticTensorDictSequential(
    net,
    extractor,
    ProbabilisticTensorDictModule(
        in_keys=["loc", "scale"],
        out_keys=["action"],
        distribution_class=dist.Normal,
        return_log_prob=True,
    ),
)
print(f"TensorDict before going through module: {td}")
td_module(td)
print(f"TensorDict after going through module now as keys action, loc and scale: {td}")
```

--------------------------------

### Install TensorDict from Source

Source: https://docs.pytorch.org/tensordict/stable/index

Installs TensorDict by cloning the repository and running a local installation. This method is recommended for developers who plan to modify the library or contribute to its codebase.

```bash
cd path/to/root
git clone https://github.com/pytorch/tensordict
cd tensordict
pip install -e .

```

--------------------------------

### Create and Use Probabilistic TensorDict Module

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.nn.ProbabilisticTensorDictModule

Demonstrates the creation and usage of a ProbabilisticTensorDictModule, including setting up underlying network modules, parameter extraction, and distribution instantiation. It shows how to pass data through the module and retrieve the resulting distribution.

```python
import torch
from tensordict import TensorDict
from tensordict.nn import (
    ProbabilisticTensorDictModule,
    ProbabilisticTensorDictSequential,
    TensorDictModule,
)
from tensordict.nn.distributions import NormalParamExtractor
from tensordict.nn.functional_modules import make_functional
from torch.distributions import Normal, Independent

td = TensorDict(
    {"input": torch.randn(3, 4), "hidden": torch.randn(3, 8)}, [3]
)
net = torch.nn.GRUCell(4, 8)
module = TensorDictModule(
    net, in_keys=["input", "hidden"], out_keys=["params"]
)
normal_params = TensorDictModule(
    NormalParamExtractor(), in_keys=["params"], out_keys=["loc", "scale"]
)
def IndepNormal(**kwargs):
    return Independent(Normal(**kwargs), 1)
prob_module = ProbabilisticTensorDictModule(
    in_keys=["loc", "scale"],
    out_keys=["action"],
    distribution_class=IndepNormal,
    return_log_prob=True,
)
td_module = ProbabilisticTensorDictSequential(
    module, normal_params, prob_module
)
params = TensorDict.from_module(td_module)
with params.to_module(td_module):
    _ = td_module(td)
print(td)
with params.to_module(td_module):
    dist = td_module.get_dist(td)
print(dist)
```

--------------------------------

### TensorDict Initialization and Basic Operations (Python)

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.TensorDict

Demonstrates how to initialize a TensorDict with a source dictionary and batch size, and perform basic operations like checking shape, unsqueezing, viewing, and cloning. It also shows how to retrieve a tensor by key.

```python
>>> import torch
>>> from tensordict import TensorDict
>>> source = {'random': torch.randn(3, 4),
...     'zeros': torch.zeros(3, 4, 5)}
>>> batch_size = [3]
>>> td = TensorDict(source, batch_size=batch_size)
>>> print(td.shape)  # equivalent to td.batch_size
torch.Size([3])
>>> td_unqueeze = td.unsqueeze(-1)
>>> print(td_unqueeze.get("zeros").shape)
torch.Size([3, 1, 4, 5])
>>> print(td_unqueeze[0].shape)
torch.Size([1])
>>> print(td_unqueeze.view(-1).shape)
torch.Size([3])
>>> print((td.clone()==td).all())
True
```

--------------------------------

### Install TensorDict from Source

Source: https://docs.pytorch.org/tensordict/stable/_sources/index.rst

Installs TensorDict by cloning the repository and installing it in editable mode. This method is suitable for developers who want to modify the library's code.

```bash
cd path/to/root
git clone https://github.com/pytorch/tensordict
cd tensordict
pip install -e .
```

--------------------------------

### Install TensorDict via Pip

Source: https://docs.pytorch.org/tensordict/stable/_sources/index.rst

Installs the stable version of the TensorDict library using pip. This is the recommended method for most users.

```bash
$ pip install tensordict
```

--------------------------------

### Stacking Modules with TensorDictSequential

Source: https://docs.pytorch.org/tensordict/stable/_downloads/ba42555713e16dc342a5f3e550dec192/tensordict_module

Demonstrates how to stack PyTorch modules (e.g., linear layers and activation functions) using TensorDictSequential. This is analogous to torch.nn.Sequential but retains intermediate inputs and outputs within a TensorDict. It shows the creation of a block and its execution on a TensorDict.

```python
block0 = TensorDictSequential(linear0, relu0)

block0(tensordict)
assert "linear0" in tensordict
assert "relu0" in tensordict
```

--------------------------------

### Example: Consuming and Displaying Streamed Data

Source: https://docs.pytorch.org/tensordict/stable/_sources/tutorials/streamed_tensordict.rst

An example function `wait7hz` demonstrates the usage of `generate_numbers` and `collect_data` to process data from a stream at a specific frequency (7 Hz) for a duration of 1 second. It then prints the collected values. The example concludes with `asyncio.run(wait7hz())` to execute the asynchronous function.

```Python
async def wait7hz() -> None:
    queue = asyncio.Queue()
    generate_task = asyncio.create_task(generate_numbers(7, queue))
    collect_data_task = asyncio.create_task(collect_data(queue, timeout=1))
    values = await collect_data_task
    generate_task.cancel()
    print(values)


asyncio.run(wait7hz())
```

--------------------------------

### Example: Loading TensorDict on Meta Device

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.LazyStackedTensorDict

Shows how to load a TensorDict onto the 'meta' device, which creates empty tensors representing the structure without loading data. Also demonstrates loading with `FakeTensorMode`.

```python
import tempfile
import torch
td = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}})
with tempfile.TemporaryDirectory() as path:
    td.save(path)
    td_load_meta = TensorDict.load_memmap(path, device="meta")
    print("meta:", td_load_meta)
    from torch._subclasses import FakeTensorMode
    with FakeTensorMode():
        td_load_fake = TensorDict.load_memmap(path)
        print("fake:", td_load_fake)
```

--------------------------------

### Create a TensorDict for input

Source: https://docs.pytorch.org/tensordict/stable/tutorials/tensordict_module

Initializes a TensorDict instance with a random tensor named 'x', preparing it as input for a module. The batch size is also specified.

```python
tensordict = TensorDict(
    x=torch.randn(5, 3),
    batch_size=[5],
)

```

--------------------------------

### Install TensorDict Nightly Build

Source: https://docs.pytorch.org/tensordict/stable/_sources/index.rst

Installs the latest nightly build of TensorDict. This is useful for users who want to access the most recent features or contribute to the development of the library.

```bash
$ pip install tensordict-nightly
```

--------------------------------

### Instantiate and Use TensorDictSequential with Basic Modules

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.nn.TensorDictSequential

This example demonstrates how to create a TensorDictSequential instance by passing in TensorDictModule instances. It shows how to call the sequential module with a TensorDict input and also with keyword arguments. The output shows the modified TensorDict with new fields and the returned tuple of tensors.

```python
>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModule, TensorDictSequential
>>> torch.manual_seed(0)
>>> module = TensorDictSequential(
...     TensorDictModule(lambda x: x+1, in_keys=["x"], out_keys=["x+1"]),
...     TensorDictModule(nn.Linear(3, 4), in_keys=["x+1"], out_keys=["w*(x+1)+b"]),
... )
>>> # with tensordict input
>>> print(module(TensorDict({"x": torch.zeros(3)}, [])))
TensorDict(
    fields={
        w*(x+1)+b: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
        x+1: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        x: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> # with tensor input: returns all the output keys in the order of the modules, ie "x+1" and "w*(x+1)+b"
>>> module(x=torch.zeros(3))
(tensor([1., 1., 1.]), tensor([-0.7214, -0.8748,  0.1571, -0.1138], grad_fn=<AddBackward0>))
>>> module(torch.zeros(3))
(tensor([1., 1., 1.]), tensor([-0.7214, -0.8748,  0.1571, -0.1138], grad_fn=<AddBackward0>))
```

--------------------------------

### LazyStackedTensorDict Assignment Example

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.LazyStackedTensorDict

Illustrates assignment to a LazyStackedTensorDict using list data structures. This example shows how to assign tensors and strings, maintaining data integrity for retrieval.

```python
td = LazyStackedTensorDict(TensorDict(), TensorDict(), stack_dim=0)
td["a"] = [torch.ones(2), torch.zeros(1)]
assert td[1]["a"] == torch.zeros(1)
td["b"] = ["a string", "another string"]
assert td[1]["b"] == "another string"
```

--------------------------------

### TensorDict Creation and Initialization

Source: https://docs.pytorch.org/tensordict/stable/genindex

This section details methods for creating TensorDict objects from various sources and initializing them.

```APIDOC
## TensorDict Creation from Various Sources

### Description
Class methods for creating TensorDict instances from different data structures like dictionaries, dataclasses, or consolidated data.

### Methods
- `from_any()`: Creates a TensorDict from any compatible input.
- `from_consolidated()`: Creates a TensorDict from consolidated data.
- `from_dataclass()`: Creates a TensorDict from a Python dataclass.
- `from_dict()`: Creates a TensorDict from a Python dictionary.

### Related Classes
These methods are class methods and can be called on `LazyStackedTensorDict`, `NonTensorStack`, `PersistentTensorDict`, `TensorDict`, `TensorDictBase`, and `TensorDictParams`.
```

--------------------------------

### Create a dataset on disk with MemmapTensor and TensorDict

Source: https://docs.pytorch.org/tensordict/stable/distributed

This example illustrates the creation of a large dataset stored on disk using `MemmapTensor` for images, masks, and labels within a `TensorDict`. It also shows how to access slices of this dataset, with the option to load queried tensors directly onto a specified device.

```python
>>> dataset = TensorDict({
...      "images": MemmapTensor(50000, 480, 480, 3),
...      "masks": MemmapTensor(50000, 480, 480, 3, dtype=torch.bool),
...      "labels": MemmapTensor(50000, 1, dtype=torch.uint8),
... }, batch_size=[50000], device="cpu")
>>> idx = [1, 5020, 34572, 11200]
>>> batch = dataset[idx].clone()
TensorDict(
    fields={
        images: Tensor(torch.Size([4, 480, 480, 3]), dtype=torch.float32),
        labels: Tensor(torch.Size([4, 1]), dtype=torch.uint8),
        masks: Tensor(torch.Size([4, 480, 480, 3]), dtype=torch.bool)},
    batch_size=torch.Size([4]),
    device=cpu,
    is_shared=False)
```

--------------------------------

### Setup and Data Loading with torchvision

Source: https://docs.pytorch.org/tensordict/stable/_downloads/a4db2295c156dd86b1f4437d354c9744/tensorclass_imagenet

This code snippet sets up the environment for data loading, including device selection and defining transformations for training and validation datasets. It utilizes torchvision's ImageFolder to load data from disk and prepares raw training data for memory mapping.

```python
import os
import time
from pathlib import Path

import torch
import torch.nn as nn
import tqdm

from tensordict import MemoryMappedTensor, tensorclass
from tensordict.utils import strtobool
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

if __name__ == "__main__":
    NUM_WORKERS = int(os.environ.get("NUM_WORKERS", "4"))
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")

    ##############################################################################
    # Transforms
    # ----------
    # First we define train and val transforms that will be applied to train and
    # val examples respectively. Note that there are random components in the
    # train transform to prevent overfitting to training data over multiple
    # epochs.

    train_transform = transforms.Compose(
        [
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )

    val_transform = transforms.Compose(
        [
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )

    ##############################################################################
    # We use ``torchvision.datasets.ImageFolder`` to conveniently load and
    # transform the data from disk.

    data_dir = Path("data") / "hymenoptera_data/"

    train_data = datasets.ImageFolder(
        root=data_dir / "train", transform=train_transform
    )
    val_data = datasets.ImageFolder(root=data_dir / "val", transform=val_transform)

    ##############################################################################
    # We'll also create a dataset of the raw training data that simply resizes
    # the image to a common size and converts to tensor. We'll use this to
    # load the data into memory-mapped tensors. The random transformations
    # need to be different each time we iterate through the data, so they
    # cannot be pre-computed. We also do not scale the data yet so that we can set the
    # ``dtype`` of the memory-mapped array to ``uint8`` and save space.

    train_data_raw = datasets.ImageFolder(
        root=data_dir / "train",
        transform=transforms.Compose(
            [transforms.Resize((256, 256)), transforms.PILToTensor()]
        ),
    )

    ##############################################################################
    # Since we'll be loading our data in batches, we write a few custom transformations
    # that take advantage of this, and apply the transformations in a vectorized way.
    #
    # First a transformation that can be used for normalization.
    class InvAffine(nn.Module):
        """A custom normalization layer."""

        def __init__(self, loc, scale):
            super().__init__()
            self.loc = loc
            self.scale = scale

        def forward(self, x):
            return (x - sel
```

--------------------------------

### Get Size of TensorDict (PyTorch)

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.PersistentTensorDict

Provides information on how to retrieve the size of a TensorDict, referencing the `batch_size` attribute. It allows getting the size of a specific dimension or the full size.

```python
# See `batch_size`.
self.size(_dim : Optional[int] = None_) -> torch.Size | int
```

--------------------------------

### Get Entry Class

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.TensorDict

Returns the class of an entry within the TensorDict, potentially avoiding an expensive `isinstance` check. This method is preferred over `tensordict.get(key).shape` when `get()` might be computationally intensive.

```python
entry_class(_key : NestedKey_) → type
    
    Returns the class of an entry, possibly avoiding a call to isinstance(td.get(key), type).
    This method should be preferred to `tensordict.get(key).shape` whenever `get()` can be expensive to execute. 
```

--------------------------------

### Distributed TensorDict send and receive example

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.LazyStackedTensorDict

Demonstrates asynchronous sending (`isend`) and receiving (`irecv`) of TensorDict contents between two processes using `torch.distributed`. This example requires `gloo` backend and TCP initialization.

```python
import torch
from tensordict import TensorDict
from torch import multiprocessing as mp

def client():
    torch.distributed.init_process_group(
        "gloo",
        rank=1,
        world_size=2,
        init_method=f"tcp://localhost:10003",
    )

    td = TensorDict({
        ("a", "b"): torch.randn(2),
        "c": torch.randn(2, 3),
        "_": torch.ones(2, 1, 5),
    },
    [2],
    )
    td.isend(0)

def server(queue, return_premature=True):
    torch.distributed.init_process_group(
        "gloo",
        rank=0,
        world_size=2,
        init_method=f"tcp://localhost:10003",
    )
    td = TensorDict({
        ("a", "b"): torch.zeros(2),
        "c": torch.zeros(2, 3),
        "_": torch.zeros(2, 1, 5),
    },
    [2],
    )
    out = td.irecv(1, return_premature=return_premature)
    if return_premature:
        for fut in out:
            fut.wait()
    assert (td != 0).all()
    queue.put("yuppie")

if __name__ == "__main__":
    queue = mp.Queue(1)
    main_worker = mp.Process(
        target=server,
        args=(queue, )
        )
    secondary_worker = mp.Process(target=client)

    main_worker.start()
    secondary_worker.start()
    out = queue.get(timeout=10)
    assert out == "yuppie"
    main_worker.join()
    secondary_worker.join()
```

--------------------------------

### TensorDict Module Operation Example

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.nn.TensorDictSequential

Demonstrates the basic operation of a TensorDictModule, showing how input keys are processed to produce output keys. It also illustrates the effect of selecting specific output keys.

```python
from tensordict import TensorDict
from tensordict.nn import TensorDictModule, TensorDictSequential
import torch
mod = TensorDictModule(lambda x, y: (x+2, y+2), in_keys=["a", "b"], out_keys=["c", "d"])
mod.select_out_keys("d")
td = TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, [])
mod(td)
mod.reset_out_keys()
mod(td)
```

--------------------------------

### Create and Initialize TensorDict

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.LazyStackedTensorDict

Demonstrates how to create a TensorDict with nested structures and tensors. It shows the initialization using keyword arguments and specifies batch sizes for different levels of nesting.

```python
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.cummin(reduce=True, dim=0)
torch.return_types.cummin(...)
```

--------------------------------

### Create and apply a TensorDictModule with a Linear layer

Source: https://docs.pytorch.org/tensordict/stable/tutorials/tensordict_module

Demonstrates creating a TensorDictModule that wraps a PyTorch Linear layer. It specifies the input key ('x') and output key ('linear0'), then applies the module to the input TensorDict.

```python
linear0 = TensorDictModule(nn.Linear(3, 128), in_keys=["x"], out_keys=["linear0"])
linear0(tensordict)

assert "linear0" in tensordict

```

--------------------------------

### TensorDict Module Basic Usage

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.nn.TensorDictSequential

Demonstrates the basic usage of a TensorDict module, including creating a TensorDict, passing it to the module, and observing the updated TensorDict with added fields. It also shows how the module can handle dispatched arguments and how to reset output keys.

```python
>>> from tensordict import TensorDict
>>> import torch
>>> from tensordict.nn import TensorDictModule

>>> mod = TensorDictModule(lambda x, y: x + y, in_keys=['a', 'b'], out_keys=['d'])
>>> td = TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, [])
>>> mod(td)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

>>> mod(torch.zeros(()), torch.ones(()))
tensor(1.)

>>> # This change will occur in-place (ie the same module will be returned with an updated list of out_keys).
>>> # It can be reverted using the `TensorDictModuleBase.reset_out_keys()` method.
>>> mod.reset_out_keys()
>>> mod(TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, []))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
```

--------------------------------

### Get data pointer of TensorDict leaves

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.TensorDict

Illustrates how to retrieve the data pointers of the leaf tensors within a TensorDict. This is useful for checking if two TensorDicts share the same underlying data. The `storage` argument can be used to get the data pointer of the raw storage.

```python
from tensordict import TensorDict
td = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2])
assert (td0.data_ptr() == td.data_ptr()).all()
td0 = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2])
td1 = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2])
td = TensorDict.lazy_stack([td0, td1])
td.data_ptr()
```

--------------------------------

### ProbabilisticTensorDictModule Example with TensorDictSequential

Source: https://docs.pytorch.org/tensordict/stable/tutorials/tensordict_module

This code snippet demonstrates the usage of ProbabilisticTensorDictModule and ProbabilisticTensorDictSequential to build a probabilistic neural network. It initializes a TensorDict, a GRUCell, a NormalParamExtractor, and then chains them together using ProbabilisticTensorDictSequential. The module is then applied to the TensorDict, showing the addition of distribution parameters ('loc', 'scale') and sampled actions ('action', 'action_log_prob').

```python
from tensordict.nn import (
    ProbabilisticTensorDictModule,
    ProbabilisticTensorDictSequential,
)
from tensordict.nn.distributions import NormalParamExtractor
from torch import distributions as dist
import torch
from tensordict import TensorDict

td = TensorDict({"input": torch.randn(3, 4), "hidden": torch.randn(3, 8)}, [3])
net = torch.nn.GRUCell(4, 8)
net = TensorDictModule(net, in_keys=["input", "hidden"], out_keys=["hidden"])
extractor = NormalParamExtractor()
extractor = TensorDictModule(extractor, in_keys=["hidden"], out_keys=["loc", "scale"])
td_module = ProbabilisticTensorDictSequential(
    net,
    extractor,
    ProbabilisticTensorDictModule(
        in_keys=["loc", "scale"],
        out_keys=["action"],
        distribution_class=dist.Normal,
        return_log_prob=True,
    ),
)
print(f"TensorDict before going through module: {td}")
td_module(td)
print(f"TensorDict after going through module now as keys action, loc and scale: {td}")
```

--------------------------------

### CudaGraphModule Requirements: No Dynamic Control Flow Example

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.nn.CudaGraphModule

Provides an example of a function that will fail when wrapped by CudaGraphModule due to dynamic control flow (if-else statement based on a condition). It highlights the need to use PyTorch's functional equivalents for such operations.

```python
>>> def func(x):
...     if x.norm() > 1:
...         return x + 1
...     else:
...         return x - 1
```

--------------------------------

### TensorDict auto_batch_size_ method examples (PyTorch)

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.PersistentTensorDict

Demonstrates the `auto_batch_size_` method of TensorDict, which automatically sets the batch size based on the tensor dimensions. Examples show how to use it without arguments to infer the full batch size and with `batch_dims` to limit the inferred batch size.

```python
>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict({"a": torch.randn(3, 4, 5), "b": {"c": torch.randn(3, 4, 6)}}, batch_size=[])
>>> td.auto_batch_size_()
>>> print(td.batch_size)
torch.Size([3, 4])
>>> td.auto_batch_size_(batch_dims=1)
>>> print(td.batch_size)
torch.Size([3])
```

--------------------------------

### Create TensorDict with Input Data

Source: https://docs.pytorch.org/tensordict/stable/_downloads/ba42555713e16dc342a5f3e550dec192/tensordict_module

Initializes a TensorDict object with a random tensor named 'x' and specifies the batch size. This TensorDict will serve as the input for subsequent TensorDictModule operations.

```python
tensordict = TensorDict(
    x=torch.randn(5, 3),
    batch_size=[5],
)
```

--------------------------------

### Combine Probabilistic and Standard TensorDict Modules Sequentially

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.nn.TensorDictSequential

This example demonstrates creating a complex sequential module by combining ProbabilisticTensorDictSequential and TensorDictSequential. It involves defining modules for parameter extraction, distribution fitting, and final output prediction. The example concludes by creating a functional version of the combined module.

```python
>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import (
...     ProbabilisticTensorDictModule,
...     ProbabilisticTensorDictSequential,
...     TensorDictModule,
...     TensorDictSequential,
... )
>>> from tensordict.nn.distributions import NormalParamExtractor
>>> from tensordict.nn.functional_modules import make_functional
>>> from torch.distributions import Normal
>>> td = TensorDict({"input": torch.randn(3, 4)}, [3,])
>>> net1 = torch.nn.Linear(4, 8)
>>> module1 = TensorDictModule(net1, in_keys=["input"], out_keys=["params"])
>>> normal_params = TensorDictModule(
...      NormalParamExtractor(), in_keys=["params"], out_keys=["loc", "scale"]
...  )
>>> td_module1 = ProbabilisticTensorDictSequential(
...     module1,
...     normal_params,
...     ProbabilisticTensorDictModule(
...         in_keys=["loc", "scale"],
...         out_keys=["hidden"],
...         distribution_class=Normal,
...         return_log_prob=True,
...     )
... )
>>> module2 = torch.nn.Linear(4, 8)
>>> td_module2 = TensorDictModule(
...    module=module2, in_keys=["hidden"], out_keys=["output"]
... )
>>> td_module = TensorDictSequential(td_module1, td_module2)
>>> params = TensorDict.from_module(td_module)
>>> with params.to_module(td_module):
```

--------------------------------

### ProbabilisticTensorDictModule

Source: https://docs.pytorch.org/tensordict/stable/reference/nn

Build distributions from network outputs and get summary statistics or samples from it.

```APIDOC
## ProbabilisticTensorDictModule

### Description
A probabilistic TD Module that allows building distributions from network outputs and obtaining summary statistics or samples.

### Method
POST (conceptual, as it's a module call)

### Endpoint
N/A (This is a class for building modules)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **in_keys** (List[str]) - Keys to read from the TensorDict for distribution parameters.
- **out_keys** (List[str]) - Keys to write the distribution parameters and samples to.
- **distribution_class** (Type[torch.distributions.Distribution]) - The class of the distribution to use (e.g., Normal).
- **return_log_prob** (bool) - Whether to compute and return the log probability of the samples.

### Request Example
```python
import torch
from tensordict import TensorDict
from tensordict.nn import TensorDictModule
from tensordict.nn.distributions import NormalParamExtractor
from tensordict.nn.prototype import ProbabilisticTensorDictModule
from torch.distributions import Normal

td = TensorDict({"input": torch.randn(3, 4), "hidden": torch.randn(3, 8)}, [3])
net = torch.nn.Sequential(torch.nn.GRUCell(4, 8), NormalParamExtractor())
module = TensorDictModule(net, in_keys=["input", "hidden"], out_keys=["loc", "scale"])
prob_module = ProbabilisticTensorDictModule(
    in_keys=["loc", "scale"],
    out_keys=["sample"],
    distribution_class=Normal,
    return_log_prob=True,
)
```

### Response
#### Success Response (200)
- **sample** (Tensor) - The sampled value(s) from the distribution.
- **sample_log_prob** (Tensor) - The log probability of the sampled value(s) (if `return_log_prob` is True).

#### Response Example
```json
{
  "sample": "Tensor(torch.Size([3, 4]), dtype=torch.float32)",
  "sample_log_prob": "Tensor(torch.Size([3, 4]), dtype=torch.float32)"
}
```
```

--------------------------------

### TensorDict Methods: T

Source: https://docs.pytorch.org/tensordict/stable/genindex

Details on TensorDict methods starting with 'T', including train, transpose, trunc, and type.

```APIDOC
## POST /websites/pytorch_tensordict_stable/train

### Description
Trains a model using TensorDictParams.

### Method
POST

### Endpoint
/websites/pytorch_tensordict_stable/train

### Parameters
#### Request Body
- **params** (TensorDictParams) - Required - Parameters for training.

### Request Example
```json
{
  "params": "<TensorDictParams object>"
}
```

### Response
#### Success Response (200)
- **result** (any) - The result of the training process.

#### Response Example
```json
{
  "result": "<training result>"
}
```

## POST /websites/pytorch_tensordict_stable/transpose

### Description
Transposes a LazyStackedTensorDict.

### Method
POST

### Endpoint
/websites/pytorch_tensordict_stable/transpose

### Parameters
#### Request Body
- **input_td** (LazyStackedTensorDict) - Required - The TensorDict to transpose.

### Request Example
```json
{
  "input_td": "<LazyStackedTensorDict object>"
}
```

### Response
#### Success Response (200)
- **transposed_td** (LazyStackedTensorDict) - The transposed TensorDict.

#### Response Example
```json
{
  "transposed_td": "<transposed TensorDict object>"
}
```

## POST /websites/pytorch_tensordict_stable/trunc

### Description
Truncates a LazyStackedTensorDict.

### Method
POST

### Endpoint
/websites/pytorch_tensordict_stable/trunc

### Parameters
#### Request Body
- **input_td** (LazyStackedTensorDict) - Required - The TensorDict to truncate.

### Request Example
```json
{
  "input_td": "<LazyStackedTensorDict object>"
}
```

### Response
#### Success Response (200)
- **truncated_td** (LazyStackedTensorDict) - The truncated TensorDict.

#### Response Example
```json
{
  "truncated_td": "<truncated TensorDict object>"
}
```

## POST /websites/pytorch_tensordict_stable/trunc_

### Description
Truncates a LazyStackedTensorDict in-place.

### Method
POST

### Endpoint
/websites/pytorch_tensordict_stable/trunc_

### Parameters
#### Request Body
- **input_td** (LazyStackedTensorDict) - Required - The TensorDict to truncate.

### Request Example
```json
{
  "input_td": "<LazyStackedTensorDict object>"
}
```

### Response
#### Success Response (200)
- **modified_td** (LazyStackedTensorDict) - The modified TensorDict.

#### Response Example
```json
{
  "modified_td": "<modified TensorDict object>"
}
```

## POST /websites/pytorch_tensordict_stable/type

### Description
Returns the type of a LazyStackedTensorDict.

### Method
POST

### Endpoint
/websites/pytorch_tensordict_stable/type

### Parameters
#### Request Body
- **input_td** (LazyStackedTensorDict) - Required - The TensorDict to check the type of.

### Request Example
```json
{
  "input_td": "<LazyStackedTensorDict object>"
}
```

### Response
#### Success Response (200)
- **type** (str) - The type of the TensorDict.

#### Response Example
```json
{
  "type": "<TensorDict type>"
}
```
```

--------------------------------

### Distribute TensorDict Initialization with init_remote and from_remote_init

Source: https://docs.pytorch.org/tensordict/stable/reference/generated/tensordict.LazyStackedTensorDict

Demonstrates how to initialize and send a TensorDict's metadata and content to a remote process using `init_remote` and receive it using `from_remote_init`. This example sets up a server and client using multiprocessing and `torch.distributed` for communication. Ensure environment variables for distributed communication are set.

```python
import os
import torch
import torch.distributed as dist
from tensordict import TensorDict, MemoryMappedTensor
import multiprocessing as mp

def server(queue):
    # Set environment variables for distributed communication
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29505"

    # Initialize the distributed backend
    dist.init_process_group("gloo", rank=0, world_size=2)

    # Create a sample tensordict
    td = (
        TensorDict(
            {
                ("a", "b"): torch.ones(2),
                "c": torch.ones(2),
                ("d", "e", "f"): MemoryMappedTensor.from_tensor(torch.ones(2, 2)),
            },
            [2],
        )
        .expand(1, 2)
        .contiguous()
    )

    # Send the tensordict metadata and content to the client
    td.init_remote(dst=1)

def client(queue):
    # Set environment variables for distributed communication
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29505"

    # Initialize the distributed backend
    dist.init_process_group("gloo", rank=1, world_size=2)

    # Receive the tensordict metadata and content from the server
    received_td = TensorDict.from_remote_init(src=0)

    # Verify that the received tensordict matches the expected structure and values
    assert set(received_td.keys()) == {"a", "c", "d"}
    assert (received_td == 1).all()

    # Signal that the test has completed successfully
    queue.put("yuppie")

if __name__ == "__main__":
    queue = mp.Queue(1)

    # Create and start the server and client processes
    main_worker = mp.Process(target=server, args=(queue,))
    secondary_worker = mp.Process(target=client, args=(queue,))

    main_worker.start()
    secondary_worker.start()

    try:
        out = queue.get(timeout=10)  # Wait for the signal with a timeout
        print(out)  # Should print "yuppie"
    finally:
        queue.close()
        main_worker.join(timeout=10)
        secondary_worker.join(timeout=10)

```
Tokens:
Raw