From e8c74a6c15e5868ff2829180429453281ea7f98c Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Wed, 12 Jan 2022 19:29:58 -0600
Subject: [PATCH] avoid copy in flatten by not calling concatenate
---
arraycontext/container/traversal.py | 5 ++++-
1 file changed, 4 insertions(+), 1 deletion(-)
diff --git a/arraycontext/container/traversal.py b/arraycontext/container/traversal.py
index c9b1282..7f08c1e 100644
--- a/arraycontext/container/traversal.py
+++ b/arraycontext/container/traversal.py
@@ -631,7 +631,10 @@ def flatten(ary: ArrayOrContainerT, actx: ArrayContext) -> Any:
_flatten(ary)
- return actx.np.concatenate(result)
+ if len(result) == 1:
+ return result[0]
+ else:
+ return actx.np.concatenate(result)
def unflatten(
--
GitLab