(
lhs: EinsumOperand, rhs: EinsumOperand, batch_dims, sum_dims, left_dims, right_dims
)
| 117 | |
| 118 | |
| 119 | def einsum_matmul( |
| 120 | lhs: EinsumOperand, rhs: EinsumOperand, batch_dims, sum_dims, left_dims, right_dims |
| 121 | ): |
| 122 | lshape = lhs.shape |
| 123 | rshape = rhs.shape |
| 124 | broadcast_lhs = 0 |
| 125 | broadcast_rhs = 0 |
| 126 | for dim in batch_dims: |
| 127 | if dim not in lshape: |
| 128 | broadcast_lhs += 1 |
| 129 | lshape = lshape + (dim,) |
| 130 | if dim not in rshape: |
| 131 | broadcast_rhs += 1 |
| 132 | rshape = rshape + (dim,) |
| 133 | # always use matmul |
| 134 | for dim in sum_dims: |
| 135 | if dim not in lshape: |
| 136 | broadcast_lhs += 1 |
| 137 | lshape = lshape + (dim,) |
| 138 | if dim not in rshape: |
| 139 | broadcast_rhs += 1 |
| 140 | rshape = rshape + (dim,) |
| 141 | |
| 142 | if broadcast_lhs: |
| 143 | lhs = lhs.broadcast(lshape) |
| 144 | |
| 145 | if broadcast_rhs: |
| 146 | rhs = rhs.broadcast(rshape) |
| 147 | |
| 148 | lhs = lhs.transpose(batch_dims + left_dims + sum_dims) |
| 149 | rhs = rhs.transpose(batch_dims + sum_dims + right_dims) |
| 150 | |
| 151 | lhs = lhs.reshape(("".join(batch_dims), "".join(left_dims), "".join(sum_dims))) |
| 152 | rhs = rhs.reshape(("".join(batch_dims), "".join(sum_dims), "".join(right_dims))) |
| 153 | # matmul TODO: sometimes unnecessary |
| 154 | oup = (lhs @ rhs).reshape(batch_dims + left_dims + right_dims) |
| 155 | return oup |
| 156 | |
| 157 | |
| 158 | def einsum_infer_output_shape(shapes): |
no test coverage detected