| 22 | |
| 23 | |
| 24 | class Arange: |
| 25 | name_generator = NameGenerator() |
| 26 | |
| 27 | def __init__(self, start, end, step: int = 1, binding_name=None) -> None: |
| 28 | self._idx = np.arange(start, end, step) |
| 29 | binding_name = ( |
| 30 | self.name_generator.get("_") if binding_name is None else binding_name |
| 31 | ) |
| 32 | self._dim_map = {binding_name: self._idx} |
| 33 | self._dim_order = [binding_name] |
| 34 | |
| 35 | @classmethod |
| 36 | def make(cls, idx, dim_map, dim_order): |
| 37 | ret = Arange(0, 0) |
| 38 | ret._idx = idx |
| 39 | ret._dim_map = dim_map |
| 40 | ret._dim_order = dim_order |
| 41 | return ret |
| 42 | |
| 43 | @property |
| 44 | def dim(self): |
| 45 | return len(self._idx.shape) |
| 46 | |
| 47 | def __handle_binary__(self, func, other): |
| 48 | if is_py_value(other): |
| 49 | return Arange.make(func(self._idx, other), self._dim_map, self._dim_order) |
| 50 | else: |
| 51 | assert isinstance(other, Arange) |
| 52 | left = self._idx |
| 53 | for i in range(other.dim): |
| 54 | left = np.expand_dims(left, axis=-1) |
| 55 | right = other._idx |
| 56 | for i in range(self.dim): |
| 57 | right = np.expand_dims(right, axis=0) |
| 58 | idx = func(left, right) |
| 59 | for k in self._dim_map.keys(): |
| 60 | assert k not in other._dim_map.keys(), "Unexpected dim name collision" |
| 61 | dim_map = {k: v for k, v in self._dim_map.items()} |
| 62 | dim_map.update(other._dim_map) |
| 63 | dim_order = self._dim_order + other._dim_order |
| 64 | return Arange.make(idx, dim_map, dim_order) |
| 65 | |
| 66 | def __handle_rbinary__(self, func, other): |
| 67 | if is_py_value(other): |
| 68 | return Arange.make(func(other, self._idx), self._dim_map, self._dim_order) |
| 69 | else: |
| 70 | assert isinstance(other, Arange) |
| 71 | left = self._idx |
| 72 | for i in range(other.dim): |
| 73 | left = np.expand_dims(left, axis=-1) |
| 74 | right = other._idx |
| 75 | for i in range(self.dim): |
| 76 | right = np.expand_dims(right, axis=0) |
| 77 | idx = func(right, left) |
| 78 | for k in self._dim_map.keys(): |
| 79 | assert k not in other._dim_map.keys(), "Unexpected dim name collision" |
| 80 | dim_map = {k: v for k, v in self._dim_map.items()} |
| 81 | dim_map.update(other._dim_map) |
no test coverage detected