From f2b286a55992fa3f26a6f729c48d3adbf149128d Mon Sep 17 00:00:00 2001 From: Amir Ebrahimi Date: Tue, 24 Dec 2024 13:07:39 -0800 Subject: [PATCH] Fix up repacked_index --- src/galois/_fields/_gf2.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/galois/_fields/_gf2.py b/src/galois/_fields/_gf2.py index 418a07975..9dac69f8a 100644 --- a/src/galois/_fields/_gf2.py +++ b/src/galois/_fields/_gf2.py @@ -566,12 +566,20 @@ def set_unpacked_slice(self, index, value): unpacked = np.unpackbits(packed, axis=-1, count=self._axis_count) unpacked[post_index] = value repacked = np.packbits(unpacked.view(np.ndarray), axis=-1) - if isinstance(packed_index[0], np.ndarray): - repacked_index = np.array(range(len(packed_index[0]))).reshape(packed_index.shape), packed_index[1] - elif isinstance(packed_index[0], int): - repacked_index = 0, packed_index[1] + if isinstance(packed_index, int) or isinstance(packed_index, slice) or len(packed_index) != 2: + repacked_index = packed_index else: - repacked_index = list(range(len(packed_index[0]))), packed_index[1] + packed_row_index, packed_col_index = packed_index + if isinstance(packed_row_index, np.ndarray): + repacked_index = np.array(range(len(packed_row_index))).reshape(packed_row_index.shape), packed_col_index + elif isinstance(packed_row_index, int): + repacked_index = 0, packed_col_index + elif isinstance(packed_row_index, slice): + repacked_index = slice(0 if packed_row_index.start is not None else packed_row_index.start, + (packed_row_index.stop - packed_row_index.start) // packed_row_index.step if packed_row_index.stop is not None else packed_row_index.stop, + 1 if packed_row_index.step is not None else packed_row_index.step), packed_col_index + else: + repacked_index = list(range(len(packed_row_index))), packed_col_index self.view(np.ndarray)[packed_index] = repacked[repacked_index]