Skip to content

v0.9.0

Choose a tag to compare

View all tags
@vmoens vmoens released this 09 Jul 16:09
· 104 commits to main since this release
v0.9.0
2ad3f0e

TensorDict 0.9.0 Release Notes

Overview

TensorDict 0.9.0 introduces significant improvements in performance, new features for lazy operations, enhanced CUDA graph support, and various bug fixes. This release focuses on stability improvements and new functionality for distributed and lazy tensor operations.

🚀 New Features

Lazy Operations and Stacking

  • to_lazystack(): New method to convert TensorDict instances to lazy stacks (#1351) (a5aab97)
  • Stack name preservation: tensordict.stack now preserves names when stacking TensorDict instances (#1348) (2053031)
  • update_batch_size in where(): Enhanced where() operation now supports update_batch_size parameter (#1365) (847a86c)
  • tolist_first(): New method for converting TensorDict to list with first-level flattening (#1334) (73fe89b)

Torch Function Integration

  • torch.maximum support: Added support for torch.maximum operation in TensorDict (#1362) (85f26e4)
  • Enhanced loss functions: Added support for torch.sum, torch.mean, torch.var and loss functions (l1, smooth_l1, mse) (#1361) (17ca2ff)

CUDA Graph Enhancements

  • CudaGraphModule.state_dict(): New method to access state dictionary of CUDA graph modules (#1346) (909907b)
  • Improved device handling: Better support for CUDA graph operations on non-zero devices (#1315) (89d05a1)
  • Stream management: Enhanced stream handling for CUDA graph operations (#1314) (2fd4843)

Non-Tensor Data Support

  • NonTensorDataBase and MetaData: New base classes for handling non-tensor data in TensorDict (#1324) (8d0241d)
  • Enhanced metadata handling: Improved support for metadata operations

Copy Operations

  • TensorDict.__copy__(): New method for creating shallow copies of TensorDict instances (#1321) (b6feadd)

Distributed tensordicts

  • broadcast tensordicts: New functionality for broadcasting TensorDict instances across different shapes (#1307) (2959863)
  • remote_init with subclasses: Enhanced remote initialization support for TensorDict subclasses (#1308) (5859a2c)
  • return_early for isend: New parameter for early return in send operations (#1306) (4012767)

🐛 Bug Fixes

TensorDict Operations

  • Fixed "none"/"None" environment variable handling (#1372) (cb104a1)
  • Fixed split_size validation in TensorDict.split() (#1370) (0bb94c0)
  • Fixed update_batch_size when source is TD and destination is LTD (#1371) (c8bfda2)
  • Fixed device argument in TensorDict constructor to respect CUDA current device (#1369) (afcbcec)
  • Fixed new_ operations on NonTensorStack (#1366) (5e67c32)
  • Fixed tensor_only construction (#1364) (75e2c26)
  • Fixed missing update_batch_size in lazy stack updates (#1359) (c9f0e40)
  • Fixed context managers update when a key is in _non_tensordict (#1353) (c11a95b)
  • Fixed tensorclass __enter__ and __exit__ methods (#1352) (08abb06)

Stacking and Chunking

  • Fixed tensordict.stack forcing all names to None when no match (#1350) (c90df00)
  • Fixed chunk/split memmap index when dim!=0 (#1345) (297a514)
  • Fixed nested key iterations for lazy stacks within tensorclasses (#1344) (2e616f8)
  • Fixed leaf check in stack function (#1341) (e03c25e)
  • Fixed nested tensorclass maybe_dense_stacks (#1340) (9c8dd2d)
  • Fixed chunk of NJTs (Nested JAX Tensors) (#1339) (3477e96)

Compilation and Device Issues

  • Fixed compilation of TensorClass with non-tensor + batch-size + device (#1337) (5c98749)
  • Fixed new_* operations for Lazy stacks (#1317) (1c8be19)
  • Fixed improper name setting in __setitem__ (#1313) (1d642b0)
  • Fixed CudaGraphModule on devices that are not 0 (#1315) (89d05a1)
  • Fixed lazy stack isend early return (#1316) (8d3c470)

Memory and Performance

  • Fixed memory leak caused by _validate_value (#1310) (a36f7f9)
  • Fixed flatten operation with start=end dim (#1333) (d9972b4)
  • Fixed expansion of lazy stacks (#1331) (56b4493)
  • Fixed return_composite defaults to True only when >1 distribution (#1328) (2c73924)

Distribution and Probabilistic Modules

  • Fixed TDParams compatibility with export (#1285) (ecdde0b)
  • Fixed better list assignment in tensorclasses (#1284) (6d8119c)
  • Fixed method _is_list_tensor_compatible missing return value (#1277) (a9cc632)
  • Fixed .item() warning on tensors that require grad (#1283) (910c953)

⚡ Performance Improvements

  • Faster _get_item: Optimized item retrieval operations (#1288) (1e33a18)
  • Dedicated validation functions: Improved validation performance (#1281) (604b471)
  • tensor_only for tensorclass: Enhanced performance for tensor-only operations (#1280) (d4bc34c)
  • Second attempt at caching validation: Improved caching mechanisms (#1311) (5f26a8b)
  • Better property handling in TC: Optimized property operations in TensorClass

🔧 Setup and CI Improvements

  • Static linking: _C extension now statically linked against Python library (#1304) (af17524)
  • Better version checking: Improved version validation in smoke tests (#1303) (e84d44f)
  • Python 3.13 support: Added support for Python 3.13 nightly builds (#1279) (0eb2ad3)
  • Enhanced CI workflows: Improved continuous integration for various platforms
  • Simplified setup: Streamlined package setup process (#1286) (fffffe5)

🚨 Deprecations and Breaking Changes

Deprecated Features

  • NormalParamWrapper: Deprecated in favor of tensordict.nn.NormalParamExtractor
  • Functional modules: is_functional, make_functional, and get_functional have been removed from tensordict

Future Changes

  • List-to-stack behavior: In version 0.10.0, lists will be automatically stacked by default. A FutureWarning will be raised if lists are assigned to TensorDict without setting the appropriate context manager.

🛠️ Quality of Life Improvements

  • Simplified error handling: Better error messages and handling in TensorDictSequential execution (#1326) (1330b72)
  • Enhanced flatten operations: Made flatten operation idempotent (#1332) (49698e2)
  • Better list handling: Improved list assignment in TensorDict instances (#1282) (6ad496b)
  • Enhanced validation: Better validation functions for different data types

📦 Dependencies

  • Python: Support for Python 3.9, 3.10, 3.11, 3.12, and 3.13
  • PyTorch: Compatible with PyTorch 1.12 and upward
  • Additional: numpy, cloudpickle, packaging, importlib_metadata, orjson (for Python < 3.13)

🔗 Migration Guide

For Users Upgrading from 0.8.0

  1. Update functional module usage: If using is_functional, make_functional, or get_functional, these have been removed
  2. NormalParamWrapper replacement: Use tensordict.nn.NormalParamExtractor instead of NormalParamWrapper
  3. List handling: Consider using the new set_list_to_stack context manager for consistent list behavior

For Developers

  • The new lazy stacking features provide better memory efficiency for large datasets
  • CUDA graph support has been enhanced for better GPU performance
  • Non-tensor data handling has been improved with new base classes

🎯 Contributors

Special thanks to all contributors who made this release possible, including:

  • Vincent Moens
  • Nikolai Karpov
  • Jiahao Li
  • Faury Louis
  • Douglas Boubert
  • Albert Bou

📝 Full Changelog

For a complete list of all changes, please refer to the git log from version 0.8.0 to 0.9.0.

Assets 26
Loading