| 97 | |
| 98 | |
| 99 | def _normalize_axis( |
| 100 | ndim: int, axis: Union[int, Iterable], reverse=False |
| 101 | ) -> Union[int, list]: |
| 102 | def convert(x): |
| 103 | x_org = x |
| 104 | if x < 0: |
| 105 | x = ndim + x |
| 106 | assert ( |
| 107 | x >= 0 and x < ndim |
| 108 | ), "axis {} is out of bounds for tensor of dimension {}".format(x_org, ndim) |
| 109 | return x |
| 110 | |
| 111 | if isinstance(axis, int): |
| 112 | return convert(axis) |
| 113 | elif isinstance(axis, Iterable): |
| 114 | axis_org = axis |
| 115 | axis = list(sorted(map(convert, axis), reverse=reverse)) |
| 116 | for i in range(len(axis) - 1): |
| 117 | assert axis[i] != axis[i + 1], "axis {} contains duplicated indices".format( |
| 118 | axis_org |
| 119 | ) |
| 120 | return axis |
| 121 | raise |
| 122 | |
| 123 | |
| 124 | _opr_map = { |