From 37fd56743b4f7f6ef83a554595374054191b5d90 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Wed, 24 May 2023 17:30:54 +0200 Subject: [PATCH] Update. --- eingather.py | 41 ++++++++++++++--------------------------- 1 file changed, 14 insertions(+), 27 deletions(-) diff --git a/eingather.py b/eingather.py index b271100..588aacb 100755 --- a/eingather.py +++ b/eingather.py @@ -19,40 +19,27 @@ def eingather(op, src, *indexes): s_all = "".join([s_src] + s_indexes) shape = tuple(all_sizes[s_all.index(v)] for v in s_dst) - idx = [] - n_index = 0 - - for i in range(src.dim()): - v = s_src[i] - if v == "_": - index, s_index = indexes[n_index], s_indexes[n_index] - n_index += 1 - - sub_idx = [] - - for i in range(index.dim()): - v = s_index[i] + def do(x,s_x): + idx = [] + n_index = 0 + + for i in range(x.dim()): + v = s_x[i] + if v == "_": + idx.append(do(indexes[n_index], s_indexes[n_index])) + n_index += 1 + else: j = s_dst.index(v) a = ( - torch.arange(index.size(i)) + torch.arange(x.size(i)) .reshape((1,) * j + (-1,) + (1,) * (len(s_dst) - j - 1)) .expand(shape) ) - sub_idx.append(a) - - index = index[sub_idx] - idx.append(index) - else: - j = s_dst.index(v) - a = ( - torch.arange(src.size(i)) - .reshape((1,) * j + (-1,) + (1,) * (len(s_dst) - j - 1)) - .expand(shape) - ) - idx.append(a) + idx.append(a) - return src[idx] + return x[idx] + return do(src,s_src) ####################### -- 2.39.5