Tensordict

Minor bug fixes

[BugFix] Fix none ref in during reduction (1090) 88c86f83
[Versioning] v0.6.2 fix cdd5cb39
[BugFix] smarter check in set_interaction_type 0b3c7783
[Versioning] 0.6.2 b26bbe35
[Quality] Better use of StrEnum in set_interaction_type 477d85bc
[Minor] print_directory_tree returns a string 05c0fe7e
[Doc] Add doc on export with nested keys d64c33d8
[Feature] Better logs of key errors in assert_close 1ef1188d
[Refactor] Make _set_dispatch_td_nn_modules compatible with compile f24e3d84
[Doc] Better docstring for to_module 178dfd9d
[Feature] `intersection` for assert_close 85833929
[Quality] Better error message for incongruent lists of keys 866943c3
[BugFix] Better repr of lazy stacks e00965c8
[BugFix] calling pad with immutable sequence (1075) d3bcb6ef
[Performance] Faster to f031bf26

**Full Changelog**: https://github.com/pytorch/tensordict/compare/v0.6.1...v0.6.2

This release offers some minor bug fixes in bytes (1059), a better handling of edge cases in keys, values and items (1058), better caching of grads (1069) and a broken reference cycle in reductions during calls to export (1056).

allowing you to get rid of tensordict altogether when exporting your models:
python
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](https://pytorch.org/tensordict/main/tutorials/export.html) to learn more about this feature.

The library integration with the PT2 stack is also further improved by the introduction of [`CudaGraphModule`](https://pytorch.org/tensordict/main/reference/generated/tensordict.nn.CudaGraphModule.html),
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](https://pytorch.org/tensordict/main/tutorials/streamed_tensordict.html).

*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 <vincentmoensgmail.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
* [Performance] Faster `__setitem__` (985) by vmoens
* [Performance] Faster clone 1043 by vmoens

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
* Mxbonn made their first contribution in https://github.com/pytorch/tensordict/pull/1042
* bhack made their first contribution in https://github.com/pytorch/tensordict/pull/1049

**Full Changelog**: https://github.com/pytorch/tensordict/compare/v0.5.0...v0.6.0

Co-authored-by: Vincent Moens <vmoensmeta.com> by albertbou92

This release is packed with new features and performance improvements.

What's new

`TensorDict.consolidate`
There is now a [`TensorDict.consolidate`](https://pytorch.org/tensordict/main/reference/generated/tensordict.TensorDictBase.html#tensordict.TensorDictBase.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`](https://github.com/pytorch/tensordict/blob/main/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):
python
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:
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):
python
td = TensorDict.load(path, device=device)


New features
* [Feature,Performance] `to(device, pin_memory, num_threads)` by vmoens in https://github.com/pytorch/tensordict/pull/846
* [Feature] Allow calls to get_mode, get_mean and get_median in case mode, mean or median is not present by vmoens in https://github.com/pytorch/tensordict/pull/804
* [Feature] Arithmetic ops for tensorclass by vmoens in https://github.com/pytorch/tensordict/pull/786
* [Feature] Best attempt to densly stack sub-tds when LazyStacked TDS are passed to maybe_dense_stack by vmoens in https://github.com/pytorch/tensordict/pull/799
* [Feature] Better dtype coverage by vmoens in https://github.com/pytorch/tensordict/pull/834
* [Feature] Change default interaction types to DETERMINISTIC by vmoens in https://github.com/pytorch/tensordict/pull/825
* [Feature] DETERMINISTIC interaction mode by vmoens in https://github.com/pytorch/tensordict/pull/824
* [Feature] Expose call_on_nested to apply and named_apply by vmoens in https://github.com/pytorch/tensordict/pull/768
* [Feature] Expose stack / cat as class methods by vmoens in https://github.com/pytorch/tensordict/pull/793
* [Feature] Load tensordicts on device, incl. meta by vmoens in https://github.com/pytorch/tensordict/pull/769
* [Feature] Make Probabilistic modules aware of CompositeDistributions out_keys by vmoens in https://github.com/pytorch/tensordict/pull/810
* [Feature] Memory-mapped nested tensors by vmoens in https://github.com/pytorch/tensordict/pull/618
* [Feature] Multithreaded apply by vmoens in https://github.com/pytorch/tensordict/pull/844
* [Feature] Multithreaded pin_memory by vmoens in https://github.com/pytorch/tensordict/pull/845
* [Feature] Support for non tensor data in h5 by vmoens in https://github.com/pytorch/tensordict/pull/772
* [Feature] TensorDict.consolidate by vmoens in https://github.com/pytorch/tensordict/pull/814
* [Feature] TensorDict.numpy() by vmoens in https://github.com/pytorch/tensordict/pull/787
* [Feature] TensorDict.replace by vmoens in https://github.com/pytorch/tensordict/pull/774
* [Feature] `out` argument in apply by vmoens in https://github.com/pytorch/tensordict/pull/794
* [Feature] `to` for consolidated TDs by vmoens in https://github.com/pytorch/tensordict/pull/851
* [Feature] `zero_grad` and `requires_grad_` by vmoens in https://github.com/pytorch/tensordict/pull/901
* [Feature] add_custom_mapping and NPE refactors by vmoens in https://github.com/pytorch/tensordict/pull/910
* [Feature] construct tds with kwargs by vmoens in https://github.com/pytorch/tensordict/pull/905
* [Feature] determinstic_sample for composite dist by vmoens in https://github.com/pytorch/tensordict/pull/827
* [Feature] expand_as by vmoens in https://github.com/pytorch/tensordict/pull/792
* [Feature] flatten and unflatten as decorators by vmoens in https://github.com/pytorch/tensordict/pull/779
* [Feature] from and to_pytree by vmoens in https://github.com/pytorch/tensordict/pull/832
* [Feature] from_modules expand_identical kwarg by vmoens in https://github.com/pytorch/tensordict/pull/911
* [Feature] grad and data for tensorclasses by vmoens in https://github.com/pytorch/tensordict/pull/904
* [Feature] isfinite, isnan, isreal by vmoens in https://github.com/pytorch/tensordict/pull/829
* [Feature] map_iter by vmoens in https://github.com/pytorch/tensordict/pull/847
* [Feature] map_names for composite dists by vmoens in https://github.com/pytorch/tensordict/pull/809
* [Feature] online edition of memory mapped tensordicts by vmoens in https://github.com/pytorch/tensordict/pull/775
* [Feature] remove distutils dependency and enable 3.12 support by GaetanLepage in https://github.com/pytorch/tensordict/pull/807
* [Feature] to_namedtuple and from_namedtuple by vmoens in https://github.com/pytorch/tensordict/pull/788
* [Feature] view(dtype) by vmoens in https://github.com/pytorch/tensordict/pull/835

Performance
* [Performance] Faster getattr in TC by vmoens in https://github.com/pytorch/tensordict/pull/912
* [Performance] Faster lock_/unclock_ when sub-tds are already locked by vmoens in https://github.com/pytorch/tensordict/pull/816
* [Performance] Faster multithreaded pin_memory by vmoens in https://github.com/pytorch/tensordict/pull/919
* [Performance] Faster tensorclass by vmoens in https://github.com/pytorch/tensordict/pull/791
* [Performance] Faster tensorclass set by vmoens in https://github.com/pytorch/tensordict/pull/880
* [Performance] Faster to-module by vmoens in https://github.com/pytorch/tensordict/pull/914

Bug Fixes
* [BugFix,CI] Fix storage filename tests by vmoens in https://github.com/pytorch/tensordict/pull/850
* [BugFix] property setter in tensorclass by vmoens in https://github.com/pytorch/tensordict/pull/813
* [BugFix] Allow any tensorclass to have a data field by vmoens in https://github.com/pytorch/tensordict/pull/906
* [BugFix] Allow fake-tensor detection pass through in torch 2.0 by vmoens in https://github.com/pytorch/tensordict/pull/802
* [BugFix] Avoid collapsing NonTensorStack when calling where by vmoens in https://github.com/pytorch/tensordict/pull/837
* [BugFix] Check if the current user has write access by MateuszGuzek in https://github.com/pytorch/tensordict/pull/781
* [BugFix] Ensure dtype is preserved with autocast by vmoens in https://github.com/pytorch/tensordict/pull/773
* [BugFix] FIx non-tensor writing in modules by vmoens in https://github.com/pytorch/tensordict/pull/822
* [BugFix] Fix (keys, values) in sub by vmoens in https://github.com/pytorch/tensordict/pull/907
* [BugFix] Fix `_make_dtype_promotion` backward compat by vmoens in https://github.com/pytorch/tensordict/pull/842
* [BugFix] Fix `pad_sequence` behavior for non-tensor attributes of tensorclass by kurtamohler in https://github.com/pytorch/tensordict/pull/884
* [BugFix] Fix builds by vmoens in https://github.com/pytorch/tensordict/pull/849
* [BugFix] Fix compile + vmap by vmoens in https://github.com/pytorch/tensordict/pull/924
* [BugFix] Fix deterministic fallback when the dist has no support by vmoens in https://github.com/pytorch/tensordict/pull/830
* [BugFix] Fix device parsing in augmented funcs by vmoens in https://github.com/pytorch/tensordict/pull/770
* [BugFix] Fix empty tuple index by vmoens in https://github.com/pytorch/tensordict/pull/811
* [BugFix] Fix fallback of deterministic samples when mean is not available by vmoens in https://github.com/pytorch/tensordict/pull/828
* [BugFix] Fix functorch dim mock by vmoens in https://github.com/pytorch/tensordict/pull/777
* [BugFix] Fix gather device by vmoens in https://github.com/pytorch/tensordict/pull/815
* [BugFix] Fix h5 auto batch size by vmoens in https://github.com/pytorch/tensordict/pull/798
* [BugFix] Fix key ordering in pointwise ops by vmoens in https://github.com/pytorch/tensordict/pull/855
* [BugFix] Fix lazy stack features (where and norm) by vmoens in https://github.com/pytorch/tensordict/pull/795
* [BugFix] Fix map by vmoens in https://github.com/pytorch/tensordict/pull/862
* [BugFix] Fix map test with fork on cuda by vmoens in https://github.com/pytorch/tensordict/pull/765
* [BugFix] Fix pad_sequence for non tensors by vmoens in https://github.com/pytorch/tensordict/pull/784
* [BugFix] Fix setting non-tensors as data in NonTensorData by vmoens in https://github.com/pytorch/tensordict/pull/864
* [BugFix] Fix stack of tensorclasses (and nontensors) by vmoens in https://github.com/pytorch/tensordict/pull/820
* [BugFix] Fix storage.filename compat with torch 2.0 by vmoens in https://github.com/pytorch/tensordict/pull/803
* [BugFix] Fix tensorclass register by vmoens in https://github.com/pytorch/tensordict/pull/817
* [BugFix] Fix torch version assertion by vmoens in https://github.com/pytorch/tensordict/pull/917
* [BugFix] Fix vmap compatibility with torch<2.2 by vmoens in https://github.com/pytorch/tensordict/pull/925
* [BugFix] Fix vmap for tensorclass by vmoens in https://github.com/pytorch/tensordict/pull/778
* [BugFix] Fix wheels by vmoens in https://github.com/pytorch/tensordict/pull/856
* [BugFix] Keep stack dim name in LazyStackedTensorDict copy ops by vmoens in https://github.com/pytorch/tensordict/pull/801
* [BugFix] Read-only compatibility in MemoryMappedTensor by vmoens in https://github.com/pytorch/tensordict/pull/780
* [BugFix] Refactor map and map_iter by vmoens in https://github.com/pytorch/tensordict/pull/869
* [BugFix] Sync cuda only if initialized by vmoens in https://github.com/pytorch/tensordict/pull/767
* [BugFix] fix _expand_to_match_shape for single bool tensor by vmoens in https://github.com/pytorch/tensordict/pull/902
* [BugFix] fix construction of lazy stacks from tds by vmoens in https://github.com/pytorch/tensordict/pull/903
* [BugFix] fix tensorclass set by vmoens in https://github.com/pytorch/tensordict/pull/854
* [BugFix] remove inplace updates when using td as a decorator by vmoens in https://github.com/pytorch/tensordict/pull/796
* [BugFix] use as_subclass in Buffer by vmoens in https://github.com/pytorch/tensordict/pull/913

Refactoring and code quality
* [Quality] Better nested detection in numpy() by vmoens in https://github.com/pytorch/tensordict/pull/800
* [Quality] Better repr of keys by vmoens in https://github.com/pytorch/tensordict/pull/897
* [Quality] fix c++ binaries formatting by vmoens in https://github.com/pytorch/tensordict/pull/859
* [Quality] non_blocking_pin instead of pin_memory by vmoens in https://github.com/pytorch/tensordict/pull/915
* [Quality] zip-strict when possible by vmoens in https://github.com/pytorch/tensordict/pull/886
* [Refactor] Better tensorclass method registration by vmoens in https://github.com/pytorch/tensordict/pull/797
* [Refactor] Make all leaves in tensorclass part of `_tensordict`, except for NonTensorData by vmoens in https://github.com/pytorch/tensordict/pull/841
* [Refactor] Refactor c++ binaries location by vmoens in https://github.com/pytorch/tensordict/pull/860
* [Refactor] Refactor is_dynamo_compile imports by vmoens in https://github.com/pytorch/tensordict/pull/916
* [Refactor] Remove `_run_checks` from `__init__` by vmoens in https://github.com/pytorch/tensordict/pull/843
* [Refactor] use from_file instead of mmap+from_buffer for readonly files by vmoens in https://github.com/pytorch/tensordict/pull/808

Others
* Bump jinja2 from 3.1.3 to 3.1.4 in /docs by dependabot in https://github.com/pytorch/tensordict/pull/840
* [Benchmark] Benchmark tensorclass ops by vmoens in https://github.com/pytorch/tensordict/pull/790
* [Benchmark] Fix recursion and cache errors in benchmarks by vmoens in https://github.com/pytorch/tensordict/pull/900
* [CI] Fix nightly build by vmoens in https://github.com/pytorch/tensordict/pull/861
* [CI] Python 3.12 compatibility by kurtamohler in https://github.com/pytorch/tensordict/pull/818
* [Doc] Fix symbolic trace reference in doc by vmoens in https://github.com/pytorch/tensordict/pull/918
* [Formatting] Lint revamp by vmoens in https://github.com/pytorch/tensordict/pull/890
* [Test] Test FC of memmap save and load by vmoens in https://github.com/pytorch/tensordict/pull/838
* [Versioning] Allow any torch version for local builds by vmoens in https://github.com/pytorch/tensordict/pull/764
* [Versioning] Make dependence on uint16 optional for older PT versions by vmoens in https://github.com/pytorch/tensordict/pull/839
* [Versioning] tree_leaves for pytorch < 2.3 by vmoens in https://github.com/pytorch/tensordict/pull/806
* [Versioning] v0.5 bump by vmoens in https://github.com/pytorch/tensordict/pull/848

New Contributors
* MateuszGuzek made their first contribution in https://github.com/pytorch/tensordict/pull/781
* GaetanLepage made their first contribution in https://github.com/pytorch/tensordict/pull/807
* kurtamohler made their first contribution in https://github.com/pytorch/tensordict/pull/818

**Full Changelog**: https://github.com/pytorch/tensordict/compare/v0.4.0...v0.5.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 `+=` or `data.mul_`, fused cuda kernels will be used which will drastically improve the runtime.
See
- [Feature] Pointwise arithmetic operations using foreach by vmoens in https://github.com/pytorch/tensordict/pull/722
- [Feature] Mean, std, var, prod, sum by vmoens in https://github.com/pytorch/tensordict/pull/751

- 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 https://github.com/pytorch/tensordict/pull/718
[BugFix] consistent use of non_blocking in tensordict and torch.Tensor by vmoens in https://github.com/pytorch/tensordict/pull/734
[Feature] non_blocking=None by default by vmoens in https://github.com/pytorch/tensordict/pull/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 https://github.com/pytorch/tensordict/pull/694
[BugFix] Fix inheritance from non-tensor by vmoens in https://github.com/pytorch/tensordict/pull/709
[Feature] Allow non-tensordata to be shared across processes + memmap by vmoens in https://github.com/pytorch/tensordict/pull/699
[Feature] Better detection of non-tensor data by vmoens in https://github.com/pytorch/tensordict/pull/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 the `autocast=True` argument
[Feature] Type casting for tensorclass by vmoens in https://github.com/pytorch/tensordict/pull/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 https://github.com/pytorch/tensordict/pull/695
[Feature] map with preallocated output by vmoens in https://github.com/pytorch/tensordict/pull/667

- Miscellaneous performance improvements
[Performance] Faster flatten_keys by vmoens in https://github.com/pytorch/tensordict/pull/727
[Performance] Faster update_ by vmoens in https://github.com/pytorch/tensordict/pull/705
[Performance] Minor efficiency improvements by vmoens in https://github.com/pytorch/tensordict/pull/703
[Performance] Random speedups by albanD in https://github.com/pytorch/tensordict/pull/728
[Feature] Faster to(device) by vmoens in https://github.com/pytorch/tensordict/pull/740

- Finally, we have opened a discord channel for tensordict!
[Badge] Discord shield by vmoens in https://github.com/pytorch/tensordict/pull/736

- We cleaned up the API a bit, creating a `save` and a `load` methods, or adding some utils such as `fromkeys`. One can also check if a key belongs to a tensordict as it is done with a regular dictionary with `key in tensordict`!
[Feature] __contains__, clear and fromkeys by vmoens in https://github.com/pytorch/tensordict/pull/721

Thanks for all our contributors and community for the support!

Other PRs
* [Benchmark] Benchmark to_module by vmoens in https://github.com/pytorch/tensordict/pull/669
* [Benchmark] Benchmark update_ by vmoens in https://github.com/pytorch/tensordict/pull/704
* [BugFIx] FIx tc.update_ by vmoens in https://github.com/pytorch/tensordict/pull/750
* [BugFix, Feature] `pad_sequence` refactoring by dtsaras in https://github.com/pytorch/tensordict/pull/652
* [BugFix, Feature] tensorclass.to_dict and from_dict by vmoens in https://github.com/pytorch/tensordict/pull/707
* [BugFix, Performance] Fewer imports at root by vmoens in https://github.com/pytorch/tensordict/pull/682
* [BugFix, Refactor] More reliable Sequential.get_dist by vmoens in https://github.com/pytorch/tensordict/pull/678
* [BugFix,Feature] filter_empty in apply by vmoens in https://github.com/pytorch/tensordict/pull/661
* [BugFix] Allow device overriding with None in apply by vmoens in https://github.com/pytorch/tensordict/pull/720
* [BugFix] Avoid lazy stacks in stack if not asked explicitly by vmoens in https://github.com/pytorch/tensordict/pull/741
* [BugFix] Dense stack lazy tds defaults to dense_stack_tds by vmoens in https://github.com/pytorch/tensordict/pull/713
* [BugFix] Faster to_module by vmoens in https://github.com/pytorch/tensordict/pull/670
* [BugFix] Fix colab in tutos by vmoens in https://github.com/pytorch/tensordict/pull/757
* [BugFix] Fix dense stack usage in torch.stack by vmoens in https://github.com/pytorch/tensordict/pull/714
* [BugFix] Fix empty(recurse) call in _apply_nest by vmoens in https://github.com/pytorch/tensordict/pull/658
* [BugFix] Fix indexing (lazy stack and names) by vmoens in https://github.com/pytorch/tensordict/pull/657
* [BugFix] Fix keys for nested lazy stacks by vmoens in https://github.com/pytorch/tensordict/pull/745
* [BugFix] Fix lazy params init by vmoens in https://github.com/pytorch/tensordict/pull/681
* [BugFix] Fix lazy stack keys by vmoens in https://github.com/pytorch/tensordict/pull/744
* [BugFix] Fix load_state_dict for TensorDictParams by vmoens in https://github.com/pytorch/tensordict/pull/689
* [BugFix] Fix name gathering with tensor indices by vmoens in https://github.com/pytorch/tensordict/pull/690
* [BugFix] Fix singleton dims in expand_as_right by vmoens in https://github.com/pytorch/tensordict/pull/687
* [BugFix] Fix to_module `__exit__` update when td is locked by vmoens in https://github.com/pytorch/tensordict/pull/671
* [BugFix] Fix to_module batch-size mismatch by vmoens in https://github.com/pytorch/tensordict/pull/688
* [BugFix] Fix torch_function for uninit param by vmoens in https://github.com/pytorch/tensordict/pull/683
* [BugFix] Fix zipping in flatten_keys by vmoens in https://github.com/pytorch/tensordict/pull/729
* [BugFix] Improve update_ by vmoens in https://github.com/pytorch/tensordict/pull/655
* [BugFix] Keep dim names in transpose by vmoens in https://github.com/pytorch/tensordict/pull/662
* [BugFix] Loading phantom state-dicts by vmoens in https://github.com/pytorch/tensordict/pull/650
* [BugFix] Make functorch.dim optional by vmoens in https://github.com/pytorch/tensordict/pull/737
* [BugFix] Missing **kwargs in apply_ fallback by vmoens in https://github.com/pytorch/tensordict/pull/664
* [BugFix] Patch pad_sequence by vmoens in https://github.com/pytorch/tensordict/pull/742
* [BugFix] Remove monkey patching of uninit params by vmoens in https://github.com/pytorch/tensordict/pull/684
* [BugFix] Support empty tuple in lazy stack indexing by vmoens in https://github.com/pytorch/tensordict/pull/696
* [BugFix] Track sub-tds in memmap by vmoens in https://github.com/pytorch/tensordict/pull/719
* [BugFix] Unlock td during update in to_module by vmoens in https://github.com/pytorch/tensordict/pull/686
* [BugFix] module hook fixes by vmoens in https://github.com/pytorch/tensordict/pull/673
* [BugFix] tensorclass as a decorator by vmoens in https://github.com/pytorch/tensordict/pull/691
* [CI] Doc on release tag by vmoens in https://github.com/pytorch/tensordict/pull/761
* [CI] Fix wheels by vmoens in https://github.com/pytorch/tensordict/pull/763
* [CI] Pinning mpmath by vmoens in https://github.com/pytorch/tensordict/pull/697
* [CI] Remove OSX x86 jobs by vmoens in https://github.com/pytorch/tensordict/pull/753
* [CI] Remove all osx x86 workflows by vmoens in https://github.com/pytorch/tensordict/pull/760
* [CI] Remove snapshot from CI by vmoens in https://github.com/pytorch/tensordict/pull/701
* [CI] Schedule workflow for release branches by vmoens in https://github.com/pytorch/tensordict/pull/759
* [CI] Unpin mpmath by vmoens in https://github.com/pytorch/tensordict/pull/702
* [Doc,CI] Sanitize version by vmoens in https://github.com/pytorch/tensordict/pull/762
* [Doc] Fix EnsembleModule docstring by BY571 in https://github.com/pytorch/tensordict/pull/712
* [Doc] Fix probabilistic td module doc by vmoens in https://github.com/pytorch/tensordict/pull/756
* [Doc] Installation instructions in API ref by vmoens in https://github.com/pytorch/tensordict/pull/660
* [Doc] Per-release docs by vmoens in https://github.com/pytorch/tensordict/pull/758
* [Doc] fix typo by husisy in https://github.com/pytorch/tensordict/pull/724
* [Feature] Adds utils method isin by albertbou92 in https://github.com/pytorch/tensordict/pull/654
* [Feature] Adds utils method remove_duplicates by albertbou92 in https://github.com/pytorch/tensordict/pull/653
* [Feature] Inherit lock in shape / tensor ops by vmoens in https://github.com/pytorch/tensordict/pull/730
* [Feature] Store non tensor stacks in a single json by vmoens in https://github.com/pytorch/tensordict/pull/711
* [Feature] TensorDict logger by vmoens in https://github.com/pytorch/tensordict/pull/710
* [Feature] TensorDict.depth property by vmoens in https://github.com/pytorch/tensordict/pull/732
* [Feature] Use generator for map by vmoens in https://github.com/pytorch/tensordict/pull/672
* [Feature] Warn when `reset_parameters_recursive` is a no-op by matteobettini in https://github.com/pytorch/tensordict/pull/693
* [Feature] `from_modules` method for MOE / ensemble learning by vmoens in https://github.com/pytorch/tensordict/pull/677
* [Feature] td_flatten_with_keys by vmoens in https://github.com/pytorch/tensordict/pull/675
* [Feature] tensordict.to_padded_tensor by vmoens in https://github.com/pytorch/tensordict/pull/723
* [Feature] use return_mask as a string in pad_sequence by dtsaras in https://github.com/pytorch/tensordict/pull/739
* [Minor] Minor improvements to tensorclass by vmoens in https://github.com/pytorch/tensordict/pull/743
* [Minor] Remove double locks on data and grad by vmoens in https://github.com/pytorch/tensordict/pull/752
* [Refactor, Feature] Default to empty batch size by vmoens in https://github.com/pytorch/tensordict/pull/674
* [Refactor] Cleanup deprecation of empty td filtering in apply by vmoens in https://github.com/pytorch/tensordict/pull/665
* [Refactor] Refactor contiguous by vmoens in https://github.com/pytorch/tensordict/pull/716
* [Refactor] Set lazy_legacy to False by default by vmoens in https://github.com/pytorch/tensordict/pull/680
* [Test] Add proper tests for torch.stack with lazy stacks by vmoens in https://github.com/pytorch/tensordict/pull/715
* [Versioning] Remove deprecated features by vmoens in https://github.com/pytorch/tensordict/pull/747
* [Versioning] Torch version by vmoens in https://github.com/pytorch/tensordict/pull/749
* [Versioning] v0.4.0 by vmoens in https://github.com/pytorch/tensordict/pull/649

New Contributors
* dtsaras made their first contribution in https://github.com/pytorch/tensordict/pull/652
* BY571 made their first contribution in https://github.com/pytorch/tensordict/pull/712
* husisy made their first contribution in https://github.com/pytorch/tensordict/pull/724
* albanD made their first contribution in https://github.com/pytorch/tensordict/pull/728

**Full Changelog**: https://github.com/pytorch/tensordict/compare/v0.3.0...v0.4.0

[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/