From 7d2dfaae1d30daa7f7769f5ff6314559ef5fa23a Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Tue, 1 Jun 2021 11:11:12 -0500
Subject: [PATCH] add stack()
---
arraycontext/impl/pytato.py | 5 +++++
1 file changed, 5 insertions(+)
diff --git a/arraycontext/impl/pytato.py b/arraycontext/impl/pytato.py
index 9e0992d..c021829 100644
--- a/arraycontext/impl/pytato.py
+++ b/arraycontext/impl/pytato.py
@@ -96,6 +96,11 @@ class _PytatoFakeNumpyNamespace(BaseFakeNumpyNamespace):
import pytato as pt
return pt.amax(a)
+ def stack(self, arrays, axis=0):
+ import pytato as pt
+ from meshmode.dof_array import obj_or_dof_array_vectorize_n_args
+ return obj_or_dof_array_vectorize_n_args(pt.stack, arrays, axis)
+
class PytatoCompiledOperator:
def __init__(self, actx, pytato_program, input_spec, output_spec):
--
GitLab