CHERY-PICK: Axes Propagation

kaushikcfd · Sep 12, 2022 · a174a76 · a174a76
1 parent 1e8b820
commit a174a76
Show file tree

Hide file tree

Showing 7 changed files with 789 additions and 7 deletions.
diff --git a/pytato/__init__.py b/pytato/__init__.py
@@ -101,6 +101,7 @@ def set_debug_enabled(flag: bool) -> None:
                                 number_distributed_tags,
                                 execute_distributed_partition,
                                 )
+from pytato.transform.metadata import unify_axes_tags
 
 from pytato.partition import generate_code_for_partition
 
@@ -156,6 +157,8 @@ def set_debug_enabled(flag: bool) -> None:
 
         "generate_code_for_partition",
 
+        "unify_axes_tags",
+
         # sub-modules
         "analysis", "tags", "transform",
 

diff --git a/pytato/array.py b/pytato/array.py
@@ -80,6 +80,7 @@
 .. autofunction:: vdot
 .. autofunction:: broadcast_to
 .. autofunction:: squeeze
+.. autofunction:: expand_dims
 .. automodule:: pytato.cmath
 .. automodule:: pytato.reductions
 

diff --git a/pytato/tags.py b/pytato/tags.py
@@ -9,6 +9,7 @@
 .. autoclass:: Named
 .. autoclass:: PrefixNamed
 .. autoclass:: AssumeNonNegative
+.. autoclass:: ExpandedDimsReshape
 """
 
 from typing import Tuple