From 105cf9c8691f3317765a64966cc3240e4da601db Mon Sep 17 00:00:00 2001 From: leimao Date: Mon, 21 Oct 2024 21:55:01 -0700 Subject: [PATCH] Update --- python/pycute/int_tuple.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/python/pycute/int_tuple.py b/python/pycute/int_tuple.py index da8e2a6c29..2f7cb1c98c 100644 --- a/python/pycute/int_tuple.py +++ b/python/pycute/int_tuple.py @@ -144,7 +144,25 @@ def prefix_product(a, init=1): return init -def idx2crd(idx, shape, stride=None): +def idx2crd(idx, shape): + + if is_tuple(idx): + if is_tuple(shape): # tuple tuple tuple + assert len(idx) == len(shape) + return tuple(idx2crd(i, s) for i, s in zip(idx,shape)) + else: # tuple "int" "int" + assert False # Error + else: + if is_tuple(shape): # "int" tuple tuple + tuple_idx = [] + for i in range(len(shape)): + tuple_idx.append(idx2crd(idx // product(shape[:i]), shape[i])) + return tuple(tuple_idx) + else: # "int" "int" "int" + return idx % shape + + +def idx2crd(idx, shape, stride): if stride is None: stride = prefix_product(shape)