init

setup.py

@@ -135,7 +135,7 @@ def get_extensions():
 
     ext_modules = [
         extension(
-            "tensordict._C._tensordict",
+            "tensordict._C",
             sources,
             include_dirs=[this_dir],
             extra_compile_args=extra_compile_args,

tensordict/_C/__init__.pyi

@@ -3,14 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-def unravel_key_list(list_of_keys):
-    ...
-
-def unravel_keys(keys):
-    ...
-
-def unravel_key(key):
-    ...
-
-def _unravel_key_to_tuple(key):
-    ...
+def unravel_key_list(list_of_keys): ...
+def unravel_keys(keys): ...
+def unravel_key(key): ...
+def _unravel_key_to_tuple(key): ...

tensordict/_lazy.py

@@ -45,8 +45,8 @@
     from tensordict.utils import _ftdim_mock as ftdim
 
     _has_funcdim = False
+from tensordict._C import _unravel_key_to_tuple, unravel_key_list
 from tensordict._td import _SubTensorDict, _TensorDictKeysView, TensorDict
-from tensordict._tensordict import _unravel_key_to_tuple, unravel_key_list
 from tensordict.base import (
     _is_leaf_nontensor,
     _is_tensor_collection,

tensordict/csrc/pybind.cpp

@@ -12,7 +12,7 @@
 
 namespace py = pybind11;
 
-PYBIND11_MODULE(_tensordict, m) {
+PYBIND11_MODULE(_C, m) {
   m.def("unravel_keys", &unravel_key, py::arg("key")); // for bc compat
   m.def("unravel_key", &unravel_key, py::arg("key"));
   m.def("_unravel_key_to_tuple", &_unravel_key_to_tuple, py::arg("key"));

tensordict/nn/common.py

@@ -13,9 +13,9 @@
 
 import torch
 from cloudpickle import dumps as cloudpickle_dumps, loads as cloudpickle_loads
+from tensordict._C import _unravel_key_to_tuple, unravel_key_list
 
 from tensordict._td import is_tensor_collection, TensorDictBase
-from tensordict._tensordict import _unravel_key_to_tuple, unravel_key_list
 from tensordict.functional import make_tensordict
 from tensordict.nn.functional_modules import (
     _swap_state,

tensordict/nn/sequence.py

@@ -26,7 +26,7 @@
     FUNCTORCH_ERROR = "functorch not installed. Consider installing functorch to use this functionality."
 
 
-from tensordict._tensordict import unravel_key_list
+from tensordict._C import unravel_key_list
 from tensordict.nn.common import dispatch, TensorDictModule
 from tensordict.tensordict import LazyStackedTensorDict, TensorDictBase
 from tensordict.utils import NestedKey

tensordict/tensorclass.py

@@ -30,10 +30,10 @@
 import tensordict as tensordict_lib
 
 import torch
+from tensordict._C import _unravel_key_to_tuple
 from tensordict._lazy import LazyStackedTensorDict
 from tensordict._pytree import _register_td_node
 from tensordict._td import is_tensor_collection, NO_DEFAULT, TensorDict, TensorDictBase
-from tensordict._tensordict import _unravel_key_to_tuple
 from tensordict._torch_func import TD_HANDLED_FUNCTIONS
 from tensordict.base import (
     _ACCEPTED_CLASSES,

tensordict/utils.py

@@ -46,13 +46,13 @@
 except ImportError:
     _has_funcdim = False
 from packaging.version import parse
-from tensordict._contextlib import _DecoratorContextManager
-from tensordict._tensordict import (  # noqa: F401
+from tensordict._C import (  # noqa: F401
     _unravel_key_to_tuple,
     unravel_key,
     unravel_key_list,
     unravel_keys,
 )
+from tensordict._contextlib import _DecoratorContextManager
 
 from torch import Tensor
 from torch._C import _disabled_torch_function_impl

test/test_nn.py

@@ -20,7 +20,7 @@
     tensorclass,
     TensorDict,
 )
-from tensordict._tensordict import unravel_key_list
+from tensordict._C import unravel_key_list
 from tensordict.nn import (
     dispatch,
     probabilistic as nn_probabilistic,

test/test_utils.py

@@ -10,7 +10,7 @@
 import torch
 from _utils_internal import get_available_devices
 from tensordict import TensorDict, unravel_key, unravel_key_list
-from tensordict._tensordict import _unravel_key_to_tuple
+from tensordict._C import _unravel_key_to_tuple
 from tensordict.utils import (
     _getitem_batch_size,
     _make_cache_key,