(equation, inputs: Tuple[Tensor, ...])
| 459 | |
| 460 | |
| 461 | def einsum_subgraph(equation, inputs: Tuple[Tensor, ...]) -> Tensor: |
| 462 | dtype = inputs[0].dtype |
| 463 | device = inputs[0].device |
| 464 | # assume all inputs has same dtype and device |
| 465 | for input in inputs[1:]: |
| 466 | assert input.dtype == dtype |
| 467 | assert input.device == device |
| 468 | einsum = _get_einsum_op( |
| 469 | equation, dtype, device, tuple(map(lambda x: x.ndim, inputs)) |
| 470 | ) |
| 471 | return einsum(*inputs)[0] |
| 472 | |
| 473 | |
| 474 | def einsum(equation: str, *args: Tensor) -> Tensor: |
no test coverage detected