with_tagged_axis: accept a collection of tags

inducer · Sep 24, 2024 · 7ea4699 · 7ea4699
1 parent 72ce101
commit 7ea4699
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/pytato/array.py b/pytato/array.py
@@ -714,7 +714,7 @@ def any(self, axis: int = 0) -> ArrayOrScalar:
         return pt.any(self, axis)
 
     def with_tagged_axis(self, iaxis: int,
-                         tags: Sequence[Tag] | Tag) -> Array:
+                         tags: Collection[Tag] | Tag) -> Array:
         """
         Returns a copy of *self* with *iaxis*-th axis tagged with *tags*.
         """

diff --git a/pytato/function.py b/pytato/function.py
@@ -62,11 +62,11 @@
     Any,
     Callable,
     ClassVar,
+    Collection,
     Hashable,
     Iterable,
     Iterator,
     Mapping,
-    Sequence,
     Tuple,
     TypeVar,
 )
@@ -264,7 +264,7 @@ class NamedCallResult(NamedArray):
     _mapper_method: ClassVar[str] = "map_named_call_result"
 
     def with_tagged_axis(self, iaxis: int,
-                         tags: Sequence[Tag] | Tag) -> Array:
+                         tags: Collection[Tag] | Tag) -> Array:
         raise ValueError("Tagging a NamedCallResult's axis is illegal, use"
                          " Call.with_tagged_axis instead")
 

diff --git a/pytato/transform/metadata.py b/pytato/transform/metadata.py
@@ -43,6 +43,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Collection,
     Hashable,
     Iterable,
     List,
@@ -594,11 +595,11 @@ class AxisTagAttacher(CopyMapper):
     A mapper that tags the axes in a DAG as prescribed by *axis_to_tags*.
     """
     def __init__(self,
-                 axis_to_tags: Mapping[tuple[Array, int], Iterable[Tag]],
+                 axis_to_tags: Mapping[tuple[Array, int], Collection[Tag]],
                  tag_corresponding_redn_descr: bool,
                  _function_cache: dict[Hashable, FunctionDefinition] | None = None):
         super().__init__(_function_cache=_function_cache)
-        self.axis_to_tags: Mapping[tuple[Array, int], Iterable[Tag]] = axis_to_tags
+        self.axis_to_tags: Mapping[tuple[Array, int], Collection[Tag]] = axis_to_tags
         self.tag_corresponding_redn_descr: bool = tag_corresponding_redn_descr
 
     def rec(self, expr: ArrayOrNames) -> Any: