From b964a1b8757ad1d092a49abdd013689ae9b8f128 Mon Sep 17 00:00:00 2001
From: Kaushik Kulkarni <kaushikcfd@gmail.com>
Date: Sat, 12 Feb 2022 12:15:03 -0600
Subject: [PATCH] PyOpenCLArrayContext.make_einsum: use `tagged` tags for the
output array as well
---
arraycontext/context.py | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/arraycontext/context.py b/arraycontext/context.py
index 0c9232d..2bff0ac 100644
--- a/arraycontext/context.py
+++ b/arraycontext/context.py
@@ -328,9 +328,10 @@ class ArrayContext(ABC):
arg_names = tuple("arg%d" % i for i in range(len(args)))
prg = self._get_einsum_prg(spec, arg_names, tagged)
- return self.call_loopy(
+ out_ary = self.call_loopy(
prg, **{arg_names[i]: arg for i, arg in enumerate(args)}
)["out"]
+ return self.tag(tagged, out_ary)
@abstractmethod
def clone(self):
--
GitLab