functorch does not play well with TensorDictModules (especially when callingvmap).
Therefore we need to functionalize the modules inside the TensorDictModule, and then build some interface between the TensorDictModule and the functional module.
For instance, TDSequence has to split the parameters according to each of the submodules etc.
We should test all of these functionalities:
functorch does not play well with TensorDictModules (especially when calling
vmap).Therefore we need to functionalize the modules inside the
TensorDictModule, and then build some interface between theTensorDictModuleand the functional module.For instance,
TDSequencehas to split the parameters according to each of the submodules etc.We should test all of these functionalities:
FunctionalModuleWithBuffersfrom the TDModule OR providingFunctionalModuleWithBuffersto a TDModule upon constructionProbabilisticTensorDictModule): param length, casting etc.FunctionalModuleandFunctionalModuleWithBuffers(some logics might break with the former)