diff --git a/arraycontext/impl/pytato/__init__.py b/arraycontext/impl/pytato/__init__.py index 108079a975bc515a2880edf1f4ae47fb092ff5f5..0e12b92f7b3fddcea5c869afa5277c41774f6281 100644 --- a/arraycontext/impl/pytato/__init__.py +++ b/arraycontext/impl/pytato/__init__.py @@ -219,19 +219,33 @@ class PytatoPyOpenCLArrayContext(ArrayContext): def einsum(self, spec, *args, arg_names=None, tagged=()): import pyopencl.array as cla import pytato as pt - if arg_names is not None: - from warnings import warn - warn("'arg_names' don't bear any significance in " - "PytatoPyOpenCLArrayContext.", stacklevel=2) + if arg_names is None: + arg_names = (None,) * len(args) - def preprocess_arg(arg): + def preprocess_arg(name, arg): if isinstance(arg, cla.Array): - return self.thaw(arg) + ary = self.thaw(arg) else: assert isinstance(arg, pt.Array) - return arg + ary = arg - return pt.einsum(spec, *(preprocess_arg(arg) for arg in args)) + if name is not None: + from pytato.tags import PrefixNamed + + # Tagging Placeholders with naming-related tags is pointless: + # They already have names. It's also counterproductive, as + # multiple placeholders with the same name that are not + # also the same object are not allowed, and this would produce + # a different Placeholder object of the same name. + if not isinstance(ary, pt.Placeholder): + ary = ary.tagged(PrefixNamed(name)) + + return ary + + return pt.einsum(spec, *[ + preprocess_arg(name, arg) + for name, arg in zip(arg_names, args) + ]) @property def permits_inplace_modification(self):