tag_axes: allow containers

Author: inducer

arraycontext/context.py

@@ -649,7 +649,7 @@ def permits_advanced_indexing(self) -> bool:
 def tag_axes(
         actx: ArrayContext,
         dim_to_tags: Mapping[int, ToTagSetConvertible],
-        ary: ArrayT) -> ArrayT:
+        ary: ArrayOrContainerT) -> ArrayOrContainerT:
     """
     Return a copy of *ary* with the axes in *dim_to_tags* tagged with their
     corresponding tags. Equivalent to repeated application of