| 30 | |
| 31 | |
| 32 | class EinsumOperand: |
| 33 | def __init__(self, shape: Optional[EinsumShape], tracer, ctx: EinsumContext): |
| 34 | self._shape = shape |
| 35 | self._tracer = tracer |
| 36 | self._ctx = ctx |
| 37 | |
| 38 | @property |
| 39 | def shape(self) -> EinsumShape: |
| 40 | assert self._shape is not None |
| 41 | return self._shape |
| 42 | |
| 43 | def reshape(self, shape: EinsumShape) -> "EinsumOperand": |
| 44 | if shape == self._shape: |
| 45 | return self |
| 46 | shape_val = tuple(map(lambda x: self._ctx.dims2val(x), shape)) |
| 47 | return EinsumOperand( |
| 48 | shape, self._ctx.reshape(self._tracer, shape_val), self._ctx |
| 49 | ) |
| 50 | |
| 51 | def broadcast(self, shape: EinsumShape) -> "EinsumOperand": |
| 52 | if shape == self._shape: |
| 53 | return self |
| 54 | for dim in shape: |
| 55 | assert len(dim) == 1 |
| 56 | shape_val = tuple(map(lambda x: self._ctx.dims2val(x), shape)) |
| 57 | return EinsumOperand( |
| 58 | shape, self._ctx.broadcast(self._tracer, shape_val), self._ctx |
| 59 | ) |
| 60 | |
| 61 | def transpose(self, shape: EinsumShape) -> "EinsumOperand": |
| 62 | if shape == self._shape: |
| 63 | return self |
| 64 | for dim in shape: |
| 65 | assert len(dim) == 1 |
| 66 | assert len(shape) == len(self.shape) |
| 67 | axis = [*map(lambda x: self.shape.index(x), shape)] |
| 68 | return EinsumOperand(shape, self._ctx.transpose(self._tracer, axis), self._ctx) |
| 69 | |
| 70 | def reduce(self, shape: EinsumShape) -> "EinsumOperand": |
| 71 | if shape == self._shape: |
| 72 | return self |
| 73 | for dim in shape: |
| 74 | assert len(dim) == 1 |
| 75 | assert dim in self.shape |
| 76 | axis = [*filter(lambda x: self.shape[x] not in shape, range(len(self.shape)))] |
| 77 | assert tuple([*filter(lambda x: x in shape, self.shape)]) == shape |
| 78 | return EinsumOperand(shape, self._ctx.reduce(self._tracer, axis), self._ctx) |
| 79 | |
| 80 | # TODO: Some scenarios should be re-implemented using dnn kernel |
| 81 | def diag_plane(self, shape: EinsumShape) -> "EinsumOperand": |
| 82 | # iiijkpppqqqrpppqqq -> ikjprq (may reorder) |
| 83 | if shape == self._shape: |
| 84 | return self |
| 85 | for dim in shape: |
| 86 | assert len(dim) == 1 |
| 87 | assert dim in self.shape |
| 88 | assert shape.count(dim) == 1 |
| 89 | acc = self |
no outgoing calls
no test coverage detected