From 8f77a2c06cc7c4aabc78fd93fd33253a19279a11 Mon Sep 17 00:00:00 2001
From: Kaushik Kulkarni <kaushikcfd@gmail.com>
Date: Sun, 31 Oct 2021 23:35:08 -0500
Subject: [PATCH] tag pt arrays' axes
---
arraycontext/impl/pytato/__init__.py | 8 +++-----
1 file changed, 3 insertions(+), 5 deletions(-)
diff --git a/arraycontext/impl/pytato/__init__.py b/arraycontext/impl/pytato/__init__.py
index 42619e8..f759252 100644
--- a/arraycontext/impl/pytato/__init__.py
+++ b/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
--
GitLab