From acae6c5b83a89f5de6fe375dc42c4ef2bceb270b Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Fri, 14 Jun 2024 15:12:15 +0100 Subject: [PATCH] composed map: add permute to kwarg --- pyop2/codegen/builder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyop2/codegen/builder.py b/pyop2/codegen/builder.py index 89cf31fcf..c73d2cd45 100644 --- a/pyop2/codegen/builder.py +++ b/pyop2/codegen/builder.py @@ -190,13 +190,13 @@ def __init__(self, *maps_): self.offset = maps_[0].offset self.maps_ = maps_ - def indexed(self, multiindex, layer=None): + def indexed(self, multiindex, layer=None, permute=lambda x: x): n, i, f = multiindex n_ = n for map_ in reversed(self.maps_): if map_ is not self.maps_[0]: n_, (_, _) = map_.indexed(MultiIndex(n_, FixedIndex(0), Index()), layer=None) - return self.maps_[0].indexed(MultiIndex(n_, i, f), layer=layer) + return self.maps_[0].indexed(MultiIndex(n_, i, f), layer=layer, permute=permute) class Pack(metaclass=ABCMeta):