tag pt arrays' axes

arraycontext/impl/pytato/__init__.py

@@ -212,11 +212,9 @@ class PytatoPyOpenCLArrayContext(ArrayContext):
                                        array)
 
     def tag_axis(self, iaxis, tags: Union[Sequence[Tag], Tag], array):
-        # TODO
-        from warnings import warn
-        warn("tagging PytatoPyOpenCLArrayContext's array axes: not yet implemented",
-             stacklevel=2)
-        return array
+        return rec_map_array_container(lambda x: x.with_tagged_axis(iaxis,
+                                                                     tags),
+                                        array)
 
     def einsum(self, spec, *args, arg_names=None, tagged=()):
         import pyopencl.array as cla