tensordict
TensorDict is a pytorch dedicated tensor container.
Decision gist · record as of 2026-08-14
Yes, if you work with structured tensor data in PyTorch (RL, multi-agent systems, ensemble training, or complex batched workflows). TensorDict eliminates boilerplate for keeping nested tensors synchronized and integrates well with modern PyTorch (torch.compile, vmap). Install friction is moderate due to torch dependency, but precompiled wheels and active maintenance make it straightforward. No known vulnerabilities. BSD license poses no restrictions.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires PyTorch (torch) to be installed; Python 3.10 or later.
- Medium install friction: precompiled wheels available for Python 3.10–3.13 across macOS, Linux, and Windows, but depends on torch and six other runtime packages.
- Active maintenance with a recent release (71 days old) and steady repository activity.
License · maintenance · safety
BSD (permissive) — BSD license is permissive; you may use, modify, and distribute this package with minimal restrictions, making it suitable for both open-source and commercial projects.
last release 2026-06-04 (71 days) · last repo commit 2026-08-14 · 1,034 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 2,621,852 downloads/mo, #2,967 on PyPI
Alternatives
Verify before relying
pip install tensordict
import torch
from tensordict import TensorDict
batch = TensorDict(
{"obs": torch.randn(32, 128), "action": torch.randint(0, 4, (32,))},
batch_size=[32],
)
mini = batch[:8] # slices all leaves
on_device = batch.to("cuda" if torch.cuda.is_available() else "cpu")- Whether memmap and lazy-stack features are production-ready or still experimental in 0.13.
- Performance overhead compared to hand-rolled tensor operations in typical RL or training loops.
- Compatibility with torch.compile for all documented use cases or only a subset.
What it is and what it does
TensorDict is a PyTorch container that wraps nested dictionaries of tensors and treats them as a single batched unit. Instead of manually keeping track of batch dimensions across multiple tensors or repeating operations like `.to(device)` and slicing on each leaf, you work with one TensorDict object that synchronizes all its nested tensors. Every operation—indexing, reshaping, device transfer, arithmetic—applies uniformly to every leaf tensor while maintaining a shared batch_size that validates the structure.
It is built for data-heavy machine learning workflows where the unit of computation is not a single tensor but a collection of related tensors (e.g., observations, actions, rewards, and next-state observations in reinforcement learning, or parameter sets in ensemble training). TensorDict integrates with PyTorch's ecosystem, including torch.compile for performance-critical paths, and provides utilities for memory-mapped storage, lazy stacking, and functional parameter manipulation.
Use it for
- Reinforcement learning: pass a single batch through environment, model, and loss without unpacking and repacking nested state.
- Offline dataset pipelines: store large replay buffers or training datasets with memory-mapped TensorDicts to reduce peak memory.
- Ensemble or multi-agent training: hold parameter sets or agent states as nested TensorDicts and vectorize operations over them.
- Model checkpointing: save and load structured intermediate state (activations, gradients, module parameters) in a unified format.
- Functional programming: extract module parameters into a TensorDict, swap them in and out, and compose with torch.vmap for vectorized inference.
- Data loading and preprocessing: unify heterogeneous tensor batches (images, metadata, labels) into one container that moves and reshapes atomically.
Worth the install?
AI-flagged interpretation of the facts on this page. Verify before relying on it.
Yes, if you work with structured tensor data in PyTorch (RL, multi-agent systems, ensemble training, or complex batched workflows).
TensorDict eliminates boilerplate for keeping nested tensors synchronized and integrates well with modern PyTorch (torch.compile, vmap). Install friction is moderate due to torch dependency, but precompiled wheels and active maintenance make it straightforward. No known vulnerabilities. BSD license poses no restrictions.
Install
tensordict on PyPI
Before you install
Medium install friction: precompiled wheels available for Python 3.10–3.13 across macOS, Linux, and Windows, but depends on torch and six other runtime packages. Active maintenance with a recent release (71 days old) and steady repository activity.
Requires PyTorch (torch) to be installed; Python 3.10 or later.
License in practice
BSD license is permissive; you may use, modify, and distribute this package with minimal restrictions, making it suitable for both open-source and commercial projects.
Quickstart
pip install tensordict
import torch
from tensordict import TensorDict
batch = TensorDict(
{"obs": torch.randn(32, 128), "action": torch.randint(0, 4, (32,))},
batch_size=[32],
)
mini = batch[:8] # slices all leaves
on_device = batch.to("cuda" if torch.cuda.is_available() else "cpu")Verify before relying
- Whether memmap and lazy-stack features are production-ready or still experimental in 0.13.
- Performance overhead compared to hand-rolled tensor operations in typical RL or training loops.
- Compatibility with torch.compile for all documented use cases or only a subset.
Package facts
| License | BSD permissive |
| Python support | Supports the current Python release >=3.10 |
| Install friction | Medium. Platform-specific wheel |
| Runtime dependencies | 7 packagestorchnumpycloudpicklepackagingimportlib_metadataorjsonpyvers |
| Maintenance | Actively maintained 71 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 2,621,852 / month, #2,967 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | Development Status :: 4 - BetaProgramming Language :: Python :: 3.10Programming Language :: Python :: 3.11Programming Language :: Python :: 3.12Programming Language :: Python :: 3.13Programming Language :: Python :: 3.14 |
Evidence: tensordict-0.13.0-cp310-cp310-macosx_14_0_arm64.whl; tensordict-0.13.0-cp310-cp310-manylinux_2_28_aarch64.whl; tensordict-0.13.0-cp310-cp310-manylinux_2_28_x86_64.whl; tensordict-0.13.0-cp310-cp310-win_amd64.whl; tensordict-0.13.0-cp311-cp311-macosx_14_0_arm64.whl; tensordict-0.13.0-cp311-cp311-manylinux_2_28_aarch64.whl; tensordict-0.13.0-cp311-cp311-manylinux_2_28_x86_64.whl; tensordict-0.13.0-cp311-cp311-win_amd64.whl; tensordict-0.13.0-cp312-cp312-macosx_14_0_arm64.whl; tensordict-0.13.0-cp312-cp312-manylinux_2_28_aarch64.whl; tensordict-0.13.0-cp312-cp312-manylinux_2_28_x86_64.whl; tensordict-0.13.0-cp312-cp312-win_amd64.whl; tensordict-0.13.0-cp313-cp313-macosx_14_0_arm64.whl; tensordict-0.13.0-cp313-cp313-manylinux_2_28_aarch64.whl; tensordict-0.13.0-cp313-cp313-manylinux_2_28_x86_64.whl; tensordict-0.13.0-cp313-cp313t-macosx_14_0_arm64.whl; tensordict-0.13.0-cp313-cp313t-manylinux_2_28_aarch64.whl; tensordict-0.13.0-cp313-cp313t-manylinux_2_28_x86_64.whl; tensordict-0.13.0-cp313-cp313t-win_amd64.whl; tensordict-0.13.0-cp313-cp313-win_amd64.whl
Tags
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “batched nested tensor container”
- tensordictTensorDict is a batched, nested dictionary container that behaves…
- tensordict-nightlyTensorDict is a batched, nested dictionary container for PyTorch…
- opt-einsum-fxOptimizes PyTorch einsum operations and functions containing them by…
Give your agent the search over MCP, or paste the wish link into any chat.
More Artificial Intelligence packages
LiteLLM provides a unified Python interface to call 100+ LLM providers (OpenAI, Anthropic, Gemini, Bedrock, Azure, and others) using OpenAI-compatible API format, available as both a Python SDK and a self-hosted AI Gateway proxy server.
Install it if you need to work with multiple LLM providers or want to centralize LLM routing in your organization.
Client library and CLI tool for downloading, uploading, and managing models, datasets, and repositories on the Hugging Face Hub platform.
Install it if you work with Hugging Face Hub models or datasets.
LangChain provides a framework for building agents and LLM-powered applications by composing language models, tools, and memory through a unified API that abstracts over multiple model providers.
hf-xet provides chunk-based deduplication and efficient file transfer for the Hugging Face Hub, enabling faster uploads and downloads of large files with local disk caching.
Tokenizers converts raw text into token sequences for NLP models, with support for training custom vocabularies and using pre-built tokenizers (BPE, WordPiece) optimized for speed via Rust.
Transformers provides a unified framework for loading, fine-tuning, and running state-of-the-art pretrained models across text, vision, audio, video, and multimodal tasks using PyTorch, JAX, or TensorFlow.
Install it if you need to run or train any transformer-based model for NLP, vision, audio, or multimodal tasks.
See also docarray · tensordict-nightly · tiled · objsize · tensorly · einshape · plucky · einx · einops · safetensors