add flatten_tensor_subclass to docs

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Lightning-AI · Dec 16, 2024 · 5fc2dad · 5fc2dad
1 parent 29c20ca
commit 5fc2dad
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/docs/source/reference/transforms/index.rst b/docs/source/reference/transforms/index.rst
@@ -6,5 +6,6 @@ thunder.transforms
 .. autosummary::
     :toctree: generated/
 
+    flatten_tensor_subclasses
     MaterializationTransform
     ConstantFolding
diff --git a/thunder/transforms/__init__.py b/thunder/transforms/__init__.py
@@ -1,9 +1,11 @@
 from .constant_folding import ConstantFolding
 from .materialization import MaterializationTransform
 from .qlora import LORATransform
+from .tensor_subclasses import flatten_tensor_subclasses
 
 
 __all__ = [
+    "flatten_tensor_subclasses",
     "ConstantFolding",
     "LORATransform",
     "MaterializationTransform",