Releases: pytorch/tensordict
v0.7.0: More robust composite distributions, TensorClass superclass
v0.7.0: More robust composite distributions, TensorClass superclass
v0.7.0 brings a lot of new features and bug fixes. Thanks to the vibrant community to help us keeping this project
alive!
New Contributors
A special thanks to our new contributors (+ people interacting with us on the PyTorch forum, discord or via issues and
social platforms)!
- @egaznep made their first contribution in #1056
- @priba made their first contribution in #1065
- @SamAdamDay made their first contribution in #1075
- @emmanuel-ferdman made their first contribution in #1105
- @bachdj-px made their first contribution in #1101
- @jgrhab made their first contribution in #1130
- @amdfaa made their first contribution in #1191
BC-breaking changes
- In #1180, we use the same object for
min
andmax
operations as we do with torch.Tensor.min. Previously, tensorclasses
were used, but that lead to some undefined behaviors when indexing (in PyTorch,min
returns a namedtuple that can be
indexed to get the values or argmax, whereas indexing a tensorclass indexes it along the batch dimension). - In #1166, we introduce broadcasting for pointwise operations between tensors and tensordicts. Now, the following two
operations on either sides of the==
sign are exactly equivalent:td = TensorDict(..., batch_size=[3, 5]) t = torch.randn(5) td + t == td + t.expand(td.shape)
Announced API changes
CompositeDistribution
TL;DR: We're changing the way log-probs and entropies are collected and written in ProbabilisticTensorDictModule
and
in CompositeDistribution
. The "sample_log_prob"
default key will soon be "<value>_log_prob
(or
("path", "to", "<value>_log_prob")
for nested keys). For CompositeDistribution
, a different log-prob will be
written for each leaf tensor in the distribution. This new behavior is controlled by the
tensordict.nn.set_composite_lp_aggregate(mode: bool)
function or by the COMPOSITE_LP_AGGREGATE
environment variable.
We strongly encourage users to adopt the new behavior by setting tensordict.nn.set_composite_lp_aggregate(False).set()
at the beginning of their training script.
We've had multiple rounds of refactoring for CompositeDistribution
which relied on some very specific assumptions and
resulted in a brittle and painful API. We now settled on the following API that will be enforced in v0.9, unless the
tensordict.nn.set_composite_lp_aggregate(mode)
value is explicitly set to True
(the current default).
The bulk of the problem was that log-probs were aggregated in a single tensor and registered in td["sample_log_prob"]
.
This had the following problems:
- Summing the log-probs isn't a good idea, users should be entitled to user the log-probs as they please.
"sample_log_prob"
is a generic but inappropriate name (the data may not be a randomsample
but anything else.)- Summing requires reduction (because log-probs may have different shapes), but sometimes we don't want to reduce to the
shape of the root tensordict (see pytorch/rl#2756 for instance).
What's new
tensorclass
- Introduction of the
TensorClass
class to do simple inheritance-style coding, which is accompanied by a stub file
that encodes all the TensorDict op signatures (we ensure this in the CI). See #1067 - @Tensorclass shadow attributes: you can now do
@tensorclass(shadow=True)
orclass T(TensorClass["shadow"]): ...
and you will be able to use dedicated names likeget
orvalues
as attribute names. This is slightly unsafe when you
nest the tensorclass, as we can't guarantee that the container won't be calling these methods directly on the
tensorclass. - Similarly,
@tensorclass(nocast=True)
andTensorClass["nocast"]
will deactivate the auto-casting in tensorclasses.
The behavior is now:- No value: tensorclass will cast things like TensorDict (ie,
int
ornp.arrays
will be cast totorch.Tensor
instances for example). autocast=True
will cause@tensorclass
to go one step further and attempt to cast values to the type indicated
in the dataclass definition.nocast=True
keeps values as they are. All non-tensor (or non-tensordict/tensorclass) values will be wrapped in
aNonTensorData
.
- No value: tensorclass will cast things like TensorDict (ie,
NonTensorData
- It is not easier to build non-tensor stacks through a simple
NonTensorStack(*values)
.
See the full list of features here:
[Feature] Add __abs__
docstrings, __neg__
, __rxor__
, __ror__
, __invert__
, __and__
, __rand__
, __radd__
, __rtruediv__
, __rmul__
, __rsub__
, __rpow__
, bitwise_and
, logical_and
(#1154) (d1363eb) by @vmoens ghstack-source-id: 97ce710b5a4b552d9477182e1836cf3777c2d756
[Feature] Add expln map to NormalParamExtractor (#1204) (e900b24) by @vmoens ghstack-source-id: 9003ceafbe8ecb73c701ea1ce96c0a342d0679b0
[Feature] Add missing __torch_function__
(#1169) (bc6390c) by @vmoens ghstack-source-id: 3dbefb4f5322a944664bbc2d29af7f862cb92342
[Feature] Better list casting in TensorDict.from_any (#1108) (1ffc463) by @vmoens ghstack-source-id: 427d19d5ef7c0d2779e064e64522fc0094a885af
[Feature] Better logs of key errors in assert_close (#1082) (747c593) by @vmoens ghstack-source-id: 46cb41d0da34b17ccc248119c43ddba586d29d80
[Feature] COMPOSITE_LP_AGGREGATE env variable (#1190) (9733d6e) by @vmoens ghstack-source-id: 16b07d0eac582cfd419612f87e38e1a7acffcfc0
[Feature] CompositeDistribution.from_distributions (#1113) (a45c7e3) by @vmoens ghstack-source-id: 04a62439b0fe60422fbc901172df46306e161cc5
[Feature] Ensure all dists work with DETERMINSTIC type without warning (#1182) (8e63112) by @vmoens ghstack-source-id: 63117f9b3ac4125a2be4e3e55719cc718051fc10
[Feature] Expose WrapModule (#1118) (d849756) by @vmoens ghstack-source-id: 55caa5d7c39e0f98c1e0558af2a076fee15f7984
[Feature] Fix type assertion in Seq build (#1143) (eaafc18) by @vmoens ghstack-source-id: 83d3dcafe45568c366207395a22b22fb35f61de1
[Feature] Force log_prob to return a tensordict when kwargs are passed to ProbabilisticTensorDictSequential.log_prob (#1146) (98c57ee) by @vmoens ghstack-source-id: 326d0763c9bbb13b51daac91edca4f0e821adf62
[Feature] Make ProbabilisticTensorDictSequential account for more than one distribution (#1114) (c7bd20c) by @vmoens ghstack-source-id: b62b81b5cfd49168b5875f7ba9b4f35b51cd2423
[Feature] NonTensorData(*sequence_of_any) (#1160) (70d4ed1) by @vmoens ghstack-source-id: 537f3d87b0677a1ae4992ca581a585420a10a284
[Feature] NonTensorStack.data (#1132) (4404abe) by @vmoens ghstack-source-id: 86065377cc1cd7c7283ed0a468f5d5602d60526d
[Feature] NonTensorStack.from_list (#1107) (f924afc) by @vmoens ghstack-source-id: e8f349cb06a72dcb69a639420b14406c9c08aa99
[Feature] Optional in_keys for WrapModule (#1145) (2d37d92) by @vmoens ghstack-source-id: a18dd5dff39937b027243fcebc6ef449b547e0b0
[Feature] OrderedDict for TensorDictSequential (#1142) (7df2062) by @vmoens ghstack-source-id: a8aed1eaefe066dafaa974f5b96190860de2f8f1
[Feature] ProbabilisticTensorDictModule.num_samples (#1117) (978d96c) by @vmoens ghstack-source-id: dc6b1c98cee5fefc891f0d65b66f0d17d10174ba
[Feature] ProbabilisticTensorDictSequential.default_interaction_type (#1123) (68ce9c3) by @vmoens ghstack-source-id: 37d38df36263e8accd84d6cb895269d50354e537
[Feature] Subclass conservation in td ops (#1186) (070ca61) by @vmoens ghstack-source-id: 83e79abda6a4bb6839d99240052323380981855c
[Feature] TensorClass (#1067) (a6a0dd6) by @vmoens ghstack-source-id: c3d4e17599a3204d4ad06bceb45e4fdcd0fd1be5
[Feature] TensorClass shadow attributes (#1159) (c744bcf) by @vmoens ghstack-source-id: b5cc7c7fea2d48394e63d289ee2d6f215c2333bc
[Feature] TensorDict.(dim='feature') (#1121) (ba43159) by @vmoens ghstack-source-id: 68f21aca722895e8a240dbca66e97310c20a6b5d
[Feature] TensorDict.clamp (#1165) (646683c) by @vmoens ghstack-source-id: 44f0937c195d969055de10709402af7c4473df32
[Feature] TensorDict.logsumexp (#1162) (e564b3a) by @vmoens ghstack-source-id: 84148ad9c701029db6d02dfb84ddb0a9b26c9ab7
[Feature] TensorDict.separates (#1120) (674f356) by @vmoens ghstack-source-id: be142a150bf4378a0806347257c3cf64c78e4eda
[Feature] TensorDict.softmax (#1163) (c0c6c14) by @vmoens ghstack-source-id: a88bebc23e6aaa02ec297db72dbda68ec9628ce7
[Feature] TensorDictModule in_keys allowed as Dict[str, tuple | list] to enable multi use of a sample feature (#1101) (e871b7d) by @bachdj-px
[Feature] UnbatchedTensor (#1170) (74cae09) by @vmoens ghstack-source-id: fa25726d61e913a725a71f1579eb06b09455e7c8
[Feature] intersection
for assert_close (#1078) (84d31db) by @vmoens ghstack-source-id: 3ae83c4ef90a9377405aebbf1761ace1a39417b1
[Feature] allow tensorclass to be customized (#1080) (31c7330) by @vmoens ghstack-source-id: 0b65b0a2dfb0cd7b5113e245c9444d3a0b55d085
[Feature] broadcast pointwise ops for tensor/tensordict mixed inputs (#1166) (aeff837) by @vmoens ghstack-source-id: bbefbb1a2e9841847c618bb9cf49160ff1a5c36a
[Feature] compatibility of consolidate with compile (quick version) (#1061) (3cf52a0) by @vmoens ghstack-source-id: 1bf3ca550dfe5499b58f878f72c4f1687b0f247e
[Feature] dist_params_keys and dist_sample_keys (#1179) (a728a4f) by @vmoens ghstack-source-id: d1e53e780132d04ddf37d613358b24467520230f
[Feature] flexible return type when indexing prob sequences (#1189) (790bef6) by @vmoens ghstack-source-id: 74d28ee84d965c11c527c60b20d9123ef30007f6
[Feature] from_any with UserDict (#1106) (3485c2c) by @vmoens ghstack-source-id: 420464209cff29c3a1c58ec521fbf4ed69d1355f
[Feature] inplace to method (#1066) (fbb71...
v0.6.2:
Minor bug fixes
[BugFix] Fix none ref in during reduction (#1090) 88c86f8
[Versioning] v0.6.2 fix cdd5cb3
[BugFix] smarter check in set_interaction_type 0b3c778
[Versioning] 0.6.2 b26bbe3
[Quality] Better use of StrEnum in set_interaction_type 477d85b
[Minor] print_directory_tree returns a string 05c0fe7
[Doc] Add doc on export with nested keys d64c33d
[Feature] Better logs of key errors in assert_close 1ef1188
[Refactor] Make _set_dispatch_td_nn_modules compatible with compile f24e3d8
[Doc] Better docstring for to_module 178dfd9
[Feature] intersection
for assert_close 8583392
[Quality] Better error message for incongruent lists of keys 866943c
[BugFix] Better repr of lazy stacks e00965c
[BugFix] calling pad with immutable sequence (#1075) d3bcb6e
[Performance] Faster to f031bf2
Full Changelog: v0.6.1...v0.6.2
v0.6.1: Minor fixes and perf improvements
v0.6.0: Export, streaming and `CudaGraphModule`
What's Changed
TensorDict 0.6.0 makes the @dispatch
decorator compatible with torch.export
and related APIs,
allowing you to get rid of tensordict altogether when exporting your models:
from torch.export import export
model = Seq(
# 1. A small network for embedding
Mod(nn.Linear(3, 4), in_keys=["x"], out_keys=["hidden"]),
Mod(nn.ReLU(), in_keys=["hidden"], out_keys=["hidden"]),
Mod(nn.Linear(4, 4), in_keys=["hidden"], out_keys=["latent"]),
# 2. Extracting params
Mod(NormalParamExtractor(), in_keys=["latent"], out_keys=["loc", "scale"]),
# 3. Probabilistic module
Prob(
in_keys=["loc", "scale"],
out_keys=["sample"],
distribution_class=dists.Normal,
),
)
model_export = export(model, args=(), kwargs={"x": x})
See our new tutorial to learn more about this feature.
The library integration with the PT2 stack is also further improved by the introduction of CudaGraphModule
,
which can be used to speed-up model execution under a certain set of assumptions; mainly that the inputs and outputs
are non-differentiable, that they are all tensors or constant and that the whole graph can be executed on cuda with
buffers of constant shape (ie, dynamic shape is not allowed).
We also introduce a new tutorial on streaming tensordicts.
Note: The aarch64
binaries are attached to these release notes and not available in PyPI at the moment.
Deprecations
- [Deprecate] Make calls to make_functional error #1034 by @vmoens
- [Deprecation] Act warned deprecations for v0.6 #1001 by @vmoens
- [Refactor] make TD.get default to None, like dict (#948) by @vmoens
Features
- [Feature] Allow to specify log_prob_key in CompositeDistribution (#961) by @albertbou92
- [Feature] Better typing for tensorclass #983 by @vmoens
- [Feature] Cudagraphs (#986) by @vmoens
- [Feature] Densify lazy tensordicts #955 by @vmoens
- [Feature] Frozen tensorclass #984 by @vmoens
- [Feature] Make NonTensorData a callable (#939) by @vmoens
- [Feature] NJT with lengths #1021 by @vmoens
- [Feature] Non-blocking for consolidated TD #1020 by @vmoens
- [Feature] Propagate
existsok
in memmap* methods #990 by @vmoens - [Feature] TD+NJT to(device) support #1022 by @vmoens
- [Feature] TensorDict.record_stream #1016 by @vmoens
- [Feature] Unify composite dist method signatures with other dists (#981) Co-authored-by: Vincent Moens vincentmoens@gmail.com^M
- [Feature] foreach_copy for update_ #1032 by @vmoens
- [Feature]
data_ptr()
method #1024 by @vmoens - [Feature]
inplace
arg in TDM constructor #992 by @vmoens - [Feature]
selected_out_keys
arg in TDS constructor #993 by @vmoens - [Feature] better sync and instantiation of cudagraphs (#1013) by @vmoens
- [Feature] callables for merge_tensordicts #1033 by @vmoens
- [Feature] cat and stack_from_tensordict #1018 by @vmoens
- [Feature] cat_tensors and stack_tensors #1017 by @vmoens
- [Feature] from_struct_array and to_struct_array (#938) by @vmoens
- [Feature] give a
__name__
to TDModules #1045 by @vmoens - [Feature] param_count #1046 by @vmoens
- [Feature] sorted keys, values and items #965 by @vmoens
- [Feature] str2td #953 by @vmoens
- [Feature] torch.export and onnx compatibility #991 by @vmoens
Code improvements
- [Quality] Better error for mismatching TDs (#964) by @vmoens
- [Quality] Better type hints for
__init__
(#1014) by @vmoens - [Quality] Expose private classmethods (#982) by @vmoens
- [Quality] Fewer recompiles with tensordict (#1015) by @vmoens
- [Quality] Type checks #976 by @vmoens
- [Refactor, Tests] Move TestCudagraphs by @vmoens
- [Refactor, Tests] Move TestCudagraphs #1007 by @vmoens
- [Refactor] Make @Tensorclass work properly with pyright (#1042) by @maxim
- [Refactor] Update nn inline_inbuilt check #1029 by @vmoens
- [Refactor] Use IntEnum for interaction types (#989) by @vmoens
- [Refactor] better AddStateIndependentNormalScale #1028 by @vmoens
Fixes
- [BugFix] Add nullbyte in memmap files to make fbcode happy (#943) by @vmoens
- [BugFix] Add sync to cudagraph module (#1026) by @vmoens
- [BugFix] Another compiler fix for older pytorch #980 by @vmoens
- [BugFix] Compatibility with non-tensor inputs in CudaGraphModule #1039 by @vmoens
- [BugFix] Deserializing a consolidated TD reproduces a consolidated TD #1019 by @vmoens
- [BugFix] Fix foreach_copy for older versions of PT #1035 by @vmoens
- [BugFix] Fix buffer identity in Params._apply (#1027) by @vmoens
- [BugFix] Fix key errors catch in del_ and related (#949) by @vmoens
- [BugFix] Fix number check in array parsing (np>=2 compatibility) #999 by @vmoens
- [BugFix] Fix pre 2.1 _apply compatibility #1050 by @vmoens
- [BugFix] Fix select in tensorclass (#936) by @vmoens
- [BugFix] Fix td device sync when error is raised #988 by @vmoens
- [BugFix] Fix tree_leaves import for older versions of PT #995 by @vmoens
- [BugFix] Fix vmap monkey patching #1009 by @vmoens
- [BugFix] Make probabilistic sequential modules compatible with compile #1030 by @vmoens
- [BugFix] Other dynamo fixes #977 by @vmoens
- [BugFix] Propagate maybe_dense_stack in _stack #1036 by @vmoens
- [BugFix] Regular swap_tensor for to_module in dynamo (#963) by @vmoens
- [BugFix] Remove ForkingPickler to account for change of API in torch.mp #998 by @vmoens
- [BugFix] Remove forkingpickler (#1049) by @bhack
- [BugFix] Resilient deterministic_sample for CompositeDist #1000 by @vmoens
- [BugFix] Simple syncs (#942) by @vmoens
- [BugFix] Softly revert get changes (#950) by @vmoens
- [BugFix] TDParams.to(device) works as nn.Module, not TDParams contained TD #1025 by @vmoens
- [BugFix] Use separate streams for cudagraph warmup #1010 by @vmoens
- [BugFix] dynamo compat refactors #975 by @vmoens
- [BugFix] resilient _exclude_td_from_pytree #1038 by @vmoens
- [BugFix] restrict usage of Buffers to non-batched, non-tracked tensors #979 by @vmoens
Doc
- [Doc] Broken links in GETTING_STARTED.md (#945) by @vmoens
- [Doc] Fail-on-warning in sphinx #1005 by @vmoens
- [Doc] Fix tutorials #1002 by @vmoens
- [Doc] Refactor README and add GETTING_STARTED.md (#944) by @vmoens
- [Doc] Streaming tensordicts #956 by @vmoens
- [Doc] export tutorial, TDM tuto refactoring #994 by @vmoens
Performance
Not user facing
- [Benchmark] Benchmark H2D transfer #1044 by @vmoens
- [CI, BugFix] Fix nightly build (#941) by @vmoens
- [CI] Add aarch64-linux wheels (#987) by @vmoens
- [CI] Fix versioning of h2d tests #1053 by @vmoens
- [CI] Fix windows wheels #1006 by @vmoens
- [CI] Upgrade 3.8 workflows (#967) by @vmoens
- [Minor, Format] Fix fbcode lint (#940) by @vmoens
- [Minor] Refactor is_dynamo_compiling for older torch versions (#978) by @vmoens
- [Setup] Correct read_file encoding in setup (#962) by @vmoens
- [Test] Keep a tight control over warnings (#951) by @vmoens
- [Test] Make h5py tests optional if no h5py installed (#947) by @vmoens
- [Test] Mark MP tests as slow (#946) by @vmoens
- [Test] Rename duplicated test #997 by @vmoens
- [Test] Skip compile tests that require 2.5 for stable #996 by @vmoens
- [Versioning] Versions for 0.6 (#1052) by @vmoens
New Contributors
Full Changelog: v0.5.0...v0.6.0
Co-authored-by: Vincent Moens vmoens@meta.com by @albertbou92
v0.5.0: `consolidate`, compile compatibility and better non-tensor support
This release is packed with new features and performance improvements.
What's new
TensorDict.consolidate
There is now a TensorDict.consolidate
method that will put all the tensors in a single storage. This will greatly speed-up serialization in multiprocessed and distributed settings.
PT2 support
TensorDict common ops (get
, set
, index
, arithmetic ops etc) now work within torch.compile
.
The list of supported operations can be found in test/test_compile.py
. We encourage users to report any graph break caused by tensordict to us, as we are willing to improve the coverage as much as can be.
Python 3.12 support
#807 enables python 3.12 support, a long awaited feature!
Global reduction for mean
, std
and other reduction methods
It is now possible to get the grand average of a tensordict content using tensordict.mean(reduce=True)
.
This applies to mean
, nanmean
, prod
, std
, sum
, nansum
and var
.
from_pytree
and to_pytree
We made it easy to convert a tensordict to a given pytree
structure and build it from any pytree using to_pytree
and from_pytree
. #832
Similarly, conversion to namedtuple
is now made easy thanks to #788.
map_iter
One can now iterate through a TensorDIct batch-dimension and apply a function on a separate process thanks to map_iter
.
This should enable the construction of datasets using TensorDict
, where the preproc step is executed on a separate process. #847
Using flatten and unflatten, flatten_keys and unflatten_keys as context managers
It is not possible to use flatten_keys
and flatten
as context managers (#908, #779):
with tensordict.flatten_keys() as flat_td:
flat_td["flat.key"] = 0
assert td["flat", "key"] == 0
Building a tensordict using keyword arguments
We made it easy to build tensordicts with simple keyword arguments, like a dict
is built in python:
td = TensorDict(a=0, b=1)
assert td["a"] == torch.tensor(0)
assert td["b"] == torch.tensor(1)
The batch_size
is now optional for both tensordict and tensorclasses. #905
Load tensordicts directly on device
Thanks to #769, it is now possible to load a tensordict directly on a destination device (including "meta"
device):
td = TensorDict.load(path, device=device)
New features
- [Feature,Performance]
to(device, pin_memory, num_threads)
by @vmoens in #846 - [Feature] Allow calls to get_mode, get_mean and get_median in case mode, mean or median is not present by @vmoens in #804
- [Feature] Arithmetic ops for tensorclass by @vmoens in #786
- [Feature] Best attempt to densly stack sub-tds when LazyStacked TDS are passed to maybe_dense_stack by @vmoens in #799
- [Feature] Better dtype coverage by @vmoens in #834
- [Feature] Change default interaction types to DETERMINISTIC by @vmoens in #825
- [Feature] DETERMINISTIC interaction mode by @vmoens in #824
- [Feature] Expose call_on_nested to apply and named_apply by @vmoens in #768
- [Feature] Expose stack / cat as class methods by @vmoens in #793
- [Feature] Load tensordicts on device, incl. meta by @vmoens in #769
- [Feature] Make Probabilistic modules aware of CompositeDistributions out_keys by @vmoens in #810
- [Feature] Memory-mapped nested tensors by @vmoens in #618
- [Feature] Multithreaded apply by @vmoens in #844
- [Feature] Multithreaded pin_memory by @vmoens in #845
- [Feature] Support for non tensor data in h5 by @vmoens in #772
- [Feature] TensorDict.consolidate by @vmoens in #814
- [Feature] TensorDict.numpy() by @vmoens in #787
- [Feature] TensorDict.replace by @vmoens in #774
- [Feature]
out
argument in apply by @vmoens in #794 - [Feature]
to
for consolidated TDs by @vmoens in #851 - [Feature]
zero_grad
andrequires_grad_
by @vmoens in #901 - [Feature] add_custom_mapping and NPE refactors by @vmoens in #910
- [Feature] construct tds with kwargs by @vmoens in #905
- [Feature] determinstic_sample for composite dist by @vmoens in #827
- [Feature] expand_as by @vmoens in #792
- [Feature] flatten and unflatten as decorators by @vmoens in #779
- [Feature] from and to_pytree by @vmoens in #832
- [Feature] from_modules expand_identical kwarg by @vmoens in #911
- [Feature] grad and data for tensorclasses by @vmoens in #904
- [Feature] isfinite, isnan, isreal by @vmoens in #829
- [Feature] map_iter by @vmoens in #847
- [Feature] map_names for composite dists by @vmoens in #809
- [Feature] online edition of memory mapped tensordicts by @vmoens in #775
- [Feature] remove distutils dependency and enable 3.12 support by @GaetanLepage in #807
- [Feature] to_namedtuple and from_namedtuple by @vmoens in #788
- [Feature] view(dtype) by @vmoens in #835
Performance
- [Performance] Faster getattr in TC by @vmoens in #912
- [Performance] Faster lock_/unclock_ when sub-tds are already locked by @vmoens in #816
- [Performance] Faster multithreaded pin_memory by @vmoens in #919
- [Performance] Faster tensorclass by @vmoens in #791
- [Performance] Faster tensorclass set by @vmoens in #880
- [Performance] Faster to-module by @vmoens in #914
Bug Fixes
- [BugFix,CI] Fix storage filename tests by @vmoens in #850
- [BugFix] @Property setter in tensorclass by @vmoens in #813
- [BugFix] Allow any tensorclass to have a data field by @vmoens in #906
- [BugFix] Allow fake-tensor detection pass through in torch 2.0 by @vmoens in #802
- [BugFix] Avoid collapsing NonTensorStack when calling where by @vmoens in #837
- [BugFix] Check if the current user has write access by @MateuszGuzek in #781
- [BugFix] Ensure dtype is preserved with autocast by @vmoens in #773
- [BugFix] FIx non-tensor writing in modules by @vmoens in #822
- [BugFix] Fix (keys, values) in sub by @vmoens in #907
- [BugFix] Fix
_make_dtype_promotion
backward compat by @vmoens in #842 - [BugFix] Fix
pad_sequence
behavior for non-tensor attributes of tensorclass by @kurtamohler in #884 - [BugFix] Fix builds by @vmoens in #849
- [BugFix] Fix compile + vmap by @vmoens in #924
- [BugFix] Fix deterministic fallback when the dist has no support by @vmoens in #830
- [BugFix] Fix device parsing in augmented funcs by @vmoens in #770
- [BugFix] Fix empty tuple index by @vmoens in #811
- [BugFix] Fix fallback of deterministic samples when mean is not available by @vmoens in #828
- [BugFix] Fix functorch dim mock by @vmoens in #777
- [BugFix] Fix gather device by @vmoens in #815
- [BugFix] Fix h5 auto batch size by @vmoens in #798
- [BugFix] Fix key ordering in pointwise ops by @vmoens in #855
- [BugFix] Fix lazy stack features (where and norm) by @vmoens in #795
- [BugFix] Fix map by @vmoens in #862
- [BugFix] Fix map test with fork on cuda by @vmoens in #765
- [BugFix] Fix pad_sequence for non tensors by @vmoens in #784
- [BugFix] Fix setting non-tensors as data in NonTensorData by @vmoens in https://github.com/pytorch/...
v0.4.0
What's Changed
This new version of tensordict comes with a great deal of new features:
-
You can now operate pointwise arithmetic operations on tensordict. For locked tensordicts and inplace operations such as
+=
ordata.mul_
, fused cuda kernels will be used which will drastically improve the runtime.
See -
Casting tensordicts to device is now much faster out-of-the box as data will be cast asynchronously (and it's safe too!)
[BugFix,Feature] Optional non_blocking in set, to_module and update by @vmoens in #718
[BugFix] consistent use of non_blocking in tensordict and torch.Tensor by @vmoens in #734
[Feature] non_blocking=None by default by @vmoens in #748 -
The non-tensor data API has also been improved, see
[BugFix] Allow inplace modification of non-tensor data in locked tds by @vmoens in #694
[BugFix] Fix inheritance from non-tensor by @vmoens in #709
[Feature] Allow non-tensordata to be shared across processes + memmap by @vmoens in #699
[Feature] Better detection of non-tensor data by @vmoens in #685 -
@tensorclass
now supports automatic type casting: annotating a value as a tensor or an int can ensure that the value will be cast to that type if the tensorclass decorator takes theautocast=True
argument
[Feature] Type casting for tensorclass by @vmoens in #735 -
TensorDict.map now supports the "fork" start method. Preallocated outputs are also a possibility.
[Feature] mp_start_method in tensordict map by @vmoens in #695
[Feature] map with preallocated output by @vmoens in #667 -
Miscellaneous performance improvements
[Performance] Faster flatten_keys by @vmoens in #727
[Performance] Faster update_ by @vmoens in #705
[Performance] Minor efficiency improvements by @vmoens in #703
[Performance] Random speedups by @albanD in #728
[Feature] Faster to(device) by @vmoens in #740 -
Finally, we have opened a discord channel for tensordict!
[Badge] Discord shield by @vmoens in #736 -
We cleaned up the API a bit, creating a
save
and aload
methods, or adding some utils such asfromkeys
. One can also check if a key belongs to a tensordict as it is done with a regular dictionary withkey in tensordict
!
[Feature] contains, clear and fromkeys by @vmoens in #721
Thanks for all our contributors and community for the support!
Other PRs
- [Benchmark] Benchmark to_module by @vmoens in #669
- [Benchmark] Benchmark update_ by @vmoens in #704
- [BugFIx] FIx tc.update_ by @vmoens in #750
- [BugFix, Feature]
pad_sequence
refactoring by @dtsaras in #652 - [BugFix, Feature] tensorclass.to_dict and from_dict by @vmoens in #707
- [BugFix, Performance] Fewer imports at root by @vmoens in #682
- [BugFix, Refactor] More reliable Sequential.get_dist by @vmoens in #678
- [BugFix,Feature] filter_empty in apply by @vmoens in #661
- [BugFix] Allow device overriding with None in apply by @vmoens in #720
- [BugFix] Avoid lazy stacks in stack if not asked explicitly by @vmoens in #741
- [BugFix] Dense stack lazy tds defaults to dense_stack_tds by @vmoens in #713
- [BugFix] Faster to_module by @vmoens in #670
- [BugFix] Fix colab in tutos by @vmoens in #757
- [BugFix] Fix dense stack usage in torch.stack by @vmoens in #714
- [BugFix] Fix empty(recurse) call in _apply_nest by @vmoens in #658
- [BugFix] Fix indexing (lazy stack and names) by @vmoens in #657
- [BugFix] Fix keys for nested lazy stacks by @vmoens in #745
- [BugFix] Fix lazy params init by @vmoens in #681
- [BugFix] Fix lazy stack keys by @vmoens in #744
- [BugFix] Fix load_state_dict for TensorDictParams by @vmoens in #689
- [BugFix] Fix name gathering with tensor indices by @vmoens in #690
- [BugFix] Fix singleton dims in expand_as_right by @vmoens in #687
- [BugFix] Fix to_module
__exit__
update when td is locked by @vmoens in #671 - [BugFix] Fix to_module batch-size mismatch by @vmoens in #688
- [BugFix] Fix torch_function for uninit param by @vmoens in #683
- [BugFix] Fix zipping in flatten_keys by @vmoens in #729
- [BugFix] Improve update_ by @vmoens in #655
- [BugFix] Keep dim names in transpose by @vmoens in #662
- [BugFix] Loading phantom state-dicts by @vmoens in #650
- [BugFix] Make functorch.dim optional by @vmoens in #737
- [BugFix] Missing **kwargs in apply_ fallback by @vmoens in #664
- [BugFix] Patch pad_sequence by @vmoens in #742
- [BugFix] Remove monkey patching of uninit params by @vmoens in #684
- [BugFix] Support empty tuple in lazy stack indexing by @vmoens in #696
- [BugFix] Track sub-tds in memmap by @vmoens in #719
- [BugFix] Unlock td during update in to_module by @vmoens in #686
- [BugFix] module hook fixes by @vmoens in #673
- [BugFix] tensorclass as a decorator by @vmoens in #691
- [CI] Doc on release tag by @vmoens in #761
- [CI] Fix wheels by @vmoens in #763
- [CI] Pinning mpmath by @vmoens in #697
- [CI] Remove OSX x86 jobs by @vmoens in #753
- [CI] Remove all osx x86 workflows by @vmoens in #760
- [CI] Remove snapshot from CI by @vmoens in #701
- [CI] Schedule workflow for release branches by @vmoens in #759
- [CI] Unpin mpmath by @vmoens in #702
- [Doc,CI] Sanitize version by @vmoens in #762
- [Doc] Fix EnsembleModule docstring by @BY571 in #712
- [Doc] Fix probabilistic td module doc by @vmoens in #756
- [Doc] Installation instructions in API ref by @vmoens in #660
- [Doc] Per-release docs by @vmoens in #758
- [Doc] fix typo by @husisy in #724
- [Feature] Adds utils method isin by @albertbou92 in #654
- [Feature] Adds utils method remove_duplicates by @albertbou92 in #653
- [Feature] Inherit lock in shape / tensor ops by @vmoens in #730
- [Feature] Store non tensor stacks in a single json by @vmoens in #711
- [Feature] TensorDict logger by @vmoens in #710
- [Feature] TensorDict.depth property by @vmoens in #732
- [Feature] Use generator for map by @vmoens in #672
- [Feature] Warn when
reset_parameters_recursive
is a no-op by @matteobettini in #693 - [Feature]
from_modules
method for MOE / ensemble learning by @vmoens in #677 - [Feature] td_flatten_with_keys by @vmoens in #675
- [Feature] tensordict.to_padded_tensor by @vmoens in #723
- [Feature] use return_mask as a string in pad_sequence by @dtsaras in #739
- [Minor] Minor improvements to tensorclass by @vmoens in https://github.com/pytorch/tensordict/pu...
v0.3.2: Minor release
[BugFix,Feature] Optional non_blocking in set, to_module and update (#718)
[Refactor] Refactor contiguous (#716)
[Test] Add proper tests for torch.stack with lazy stacks (#715)
[BugFix] Fix dense stack usage in torch.stack (#714)
[BugFix] Dense stack lazy tds defaults to dense_stack_tds (#713)
[Feature] Store non tensor stacks in a single json (#711)
[Feature] TensorDict logger (#710)
[BugFix, Feature] tensorclass.to_dict and from_dict (#707)
[BugFix] Fix inheritance from non-tensor (#709)
[Performance] Faster update_ (#705)
[Benchmark] Benchmark update_ (#704)
[Performance] Minor efficiency improvements (#703)
[Feature] Allow non-tensordata to be shared across processes + memmap (#699)
[CI] Unpin mpmath (#702)
[CI] Remove snapshot from CI (#701)
[BugFix] Support empty tuple in lazy stack indexing (#696)
[CI] Pinning mpmath (#697)
[BugFix] Allow inplace modification of non-tensor data in locked tds (#694)
[Feature] Better detection of non-tensor data (#685)
[Feature] Warn when reset_parameters_recursive is a no-op (#693)
[BugFix,Feature] filter_empty in apply (#661)
See the release on PyPI: https://pypi.org/project/tensordict/0.3.2/
v0.3.1
v0.3.0: `MemoryMappedTensor`, pickle-free multithreaded serialization and more!
In this release we introduce a bunch of exciting features to TensorDict:
-
We deprecate MemmapTensor in favour of MemoryMappedTensor, which is fully backed by torch.Tensor and not numpy anymore. The new API is faster and way more bug-free than it used too. See #541
-
Saving tensordicts on disk can now be done via
memmap
,memmap_
andmemmap_like
which all support multithreading. If possible, serialization is pickle free (memmap + json) andtorch.save
is only used for classes that fail to be serialized with json. Serializing models using tensordict is now 3-10x faster than using torch.save, even for SOTA LLMs such as LLAMA. -
TensorDict can now carry non tensor data through the
NonTensorData
class. Assigning non-tensor data can also be done via__setitem__
and they can be retrieved via__getitem__
. #601 -
A bunch of new operations have appeared too such as
named_apply
(apply with key names) ortensordict.auto_batch_size_()
, and operations like update can now be achieved for only a subset of keys. -
Almost all operations in the library are now faster!
-
We are slowing deprecating lazy classes except for
LazyStackedTensorDict
. Whereastorch.stack
used to systematically return a lazy stack, it now returns a dense stack if theset_lazy_legacy(mode)
decorator is set toFalse
(which will be the default in the next release). The old behaviour can be set withset_lazy_legacy(True)
. Lazy stacks can still be obtained usingLazyStackedTensorDict.lazy_stack
. Appropriate warnings are raised unless you have patched your code accordingly.
What's Changed
- [Refactor] MemoryMappedTensor by @vmoens in #541
- [Feature] Multithread memmap by @vmoens in #592
- [Refactor] Graceful as_tensor by @vmoens in #549
- [Test] Fix as_tensor test by @vmoens in #551
- Fix assignment of
str
-typed value to_device
attribute inMemmapTensor
by @kurt-stolle in #552 - [Refactor] Refactor split by @vmoens in #555
- [Refactor] Refactor implement_for by @vmoens in #556
- [Feature] Better constructors for MemoryMappedTensors by @vmoens in #557
- [CI] Fix benchmark on gpu by @vmoens in #560
- [CI] Add regular benchmarks to CI in PRs without upload by @vmoens in #561
- [Refactor] Major refactoring of codebase by @vmoens in #559
- [Benchmark] Benchmark split and chunk by @vmoens in #564
- [Performance] Faster split, chunk and unbind by @vmoens in #563
- [Feature] Consolidate functional calls by @vmoens in #565
- [Refactor] Improve functional call efficiency by @vmoens in #567
- [Refactor] Do not lock nested tensordict in tensordictparams by @vmoens in #568
- [Performance] Faster params and buffer registration in TensorDictParams by @vmoens in #569
- [BugFix] Graceful attribute error exit in TensorDictParams by @vmoens in #571
- [Refactor] Upgrade pytree import by @vmoens in #573
- [BugFix] Compatibility with missing _global_parameter_registration_hooks by @vmoens in #574
- [Feature] Seed workers in TensorDict.map by @vmoens in #562
- [Performance] Faster update by @vmoens in #572
- [Performance] Faster to_module by @vmoens in #575
- [BugFix] _FileHandler for windows by @vmoens in #577
- [Performance] Faster
__init__
by @vmoens in #576 - [Feature, Test] Add tests for partial update by @vmoens in #578
- [BugFix] No fallback on
TensorDictModule.__getattr__
for private attributes by @vmoens in #579 - [BugFix] Fix deepcopy of TensorDictParams by @vmoens in #580
- Add MANIFEST.in by @vmoens in #581
- [BugFix] Delete parameter/buffer before setting it with regular setattr in to_module by @vmoens in #583
- [Feature] named_apply and default value in apply by @vmoens in #584
- [BugFix] Faster empty_like for MemoryMappedTensor by @vmoens in #585
- [BugFix] Faster empty_like for MemoryMappedTensor (dup) by @vmoens in #586
- [BugFix] Adapt MemoryMappedTensor for torch < 2.0 by @vmoens in #587
- [Performance] Make copy_ a no-op if tensors are identical by @vmoens in #588
- [BugFix] Fix non-blocking arg in copy_ by @vmoens in #590
- [Feature] Unbind and stack tds in map with chunksize=0 by @vmoens in #589
- [Performance] Faster dispatch by @vmoens in #487
- [Feature] Saving metadata of tensorclass by @vmoens in #582
- [BugFix] Fix osx tests by @vmoens in #591
- [Feature] Weakref for unlocking tds by @vmoens in #595
- [BugFix] Fix pickling of weakrefs by @vmoens in #597
- [Feature] Return early a tensordict created through memmap with multiple threads by @vmoens in #598
- [CI] Depend on torch nightly for nightly releases by @vmoens in #599
- [Feature] Storing non-tensor data in tensordicts by @vmoens in #601
- [Feature, Test] FSDP and DTensors by @vmoens in #600
- [Minor] Fix type deletion in tensorclass load_memmap by @vmoens in #602
- [BugFix] Fix ellipsis check by @vmoens in #604
- [Feature] Best intention stack by @vmoens in #605
- [Feature] Remove and check for prints in codebase using flake8-print by @vmoens in #603
- [Doc] Doc revamp by @vmoens in #593
- [BugFix, Doc] Fix tutorial by @vmoens in #606
- [BugFix] Fix gh-pages upload by @vmoens in #607
- [BugFix] Upload content of html directly by @vmoens in #608
- [Feature] Improve in-place ops for TensorDictParams by @vmoens in #609
- [BugFix, CI] Fix GPU benchmarks by @vmoens in #611
- [Feature] inplace to_module by @vmoens in #610
- [Versioning] Bump v0.3.0 by @vmoens in #613
- [Feature] Support group specification by @lucifer1004 in #616
- [Refactor] Remove remaining MemmapTensor references by @vmoens in #617
- [Tests] Reorder and regroup tests by @vmoens in #614
- [Performance] Faster set by @vmoens in #619
- [Performance] Better shared/memmap inheritance and faster exclude by @vmoens in #621
- [Benchmark] Benchmark select, exclude and empty by @vmoens in #623
- [Feature] Improve the
empty
method by @vmoens in #622 - [BugFix] Fix is_memmap attribute for memmap_like and memmap by @vmoens in #625
- Bump jinja2 from 3.1.2 to 3.1.3 in /docs by @dependabot in #626
- [BugFix] Remove shared/memmap inheritance from clone / select / exclude by @vmoens in #624
- [BugFix] Fix
index in list
error by @vmoens in #627 - [Refactor] Make unbind call tensor.unbind by @vmoens in #628
- [Feature]
auto_batch_size_
by @vmoens in #630 - [BugFix] Fix NonTensorData interaction by @vmoens in #631
- [Doc] More doc on how to set and get non-tensor data by @vmoens in #632
- [Feature] _auto_make_functional and _dispatch_td_nn_modules by @vmoens in #633
- [BugFIx] Fix exclude indent by @vmoens in #637
- [BugFix] Limit number of threads in workers for .map() by @vmoens in #638
- [Feature] Robust to lazy_legacy set to false and context managers for reshape ops by @vmoens in #634
- [Minor] Typo in lazy legacy warnings by @Vmoe...
v0.2.1
What's Changed
- [CI] Migration CCI->GHA by @vmoens in #536
- [CI] Remove CCI runs by @vmoens in #537
- [BugFix] Fix CI by @vmoens in #538
- [Performance] Faster
to
by @vmoens in #524 - [Feature] Add
pad
argument to TensorDict.where by @vmoens in #539 - [BugFix]
torch.where
: Expand mask on the right in lazy stack tds by @vmoens in #542 - [BugFix] Better tensor allocation in memmap_like by @vmoens in #543
- [Versioning] M1 compatibility by @vmoens in #523
- [Release] v0.2.1 by @vmoens in #546
Full Changelog: v0.2.0...v0.2.1