| 402 | |
| 403 | @lru_cache(maxsize=None) |
| 404 | def _get_einsum_op( |
| 405 | equation: str, dtype, device, input_ndims |
| 406 | ) -> Callable[[Tuple[Tensor, ...]], Tensor]: |
| 407 | @subgraph_fn( |
| 408 | "Einsum", dtype=dtype, device=device, nr_inputs=len(input_ndims), gopt_level=2 |
| 409 | ) |
| 410 | def einsum(inputs, f, c): |
| 411 | from megengine.core.ops import builtin |
| 412 | |
| 413 | shapes, output_shape = einsum_parse_equation(equation, input_ndims) |
| 414 | ctx = EinsumContext() |
| 415 | dim2val_cache = {} |
| 416 | |
| 417 | def dim2val(dim: str): |
| 418 | if dim in dim2val_cache: |
| 419 | return dim2val_cache[dim] |
| 420 | for i, shape in enumerate(shapes): |
| 421 | if dim in shape: |
| 422 | axis = shape.index(dim) |
| 423 | dim2val_cache[dim] = f(builtin.GetVarShape(axis=axis), inputs[i]) |
| 424 | break |
| 425 | return dim2val_cache[dim] |
| 426 | |
| 427 | def dims2val(dims: str): |
| 428 | if len(dims) == 0: |
| 429 | return c(1, dtype=np.int32, device=device) |
| 430 | return reduce(lambda x, y: f("*", x, y), map(dim2val, dims)) |
| 431 | |
| 432 | def concat_dims(dims): |
| 433 | if len(dims) == 0: |
| 434 | return c([1], dtype=np.int32, device=device) |
| 435 | else: |
| 436 | shape = dims[0] |
| 437 | for dim in dims[1:]: |
| 438 | shape = f(builtin.Concat(axis=0, comp_node=device), shape, dim) |
| 439 | return shape |
| 440 | |
| 441 | def reduce_sum(x, axis): |
| 442 | if len(axis) == 0: |
| 443 | return x |
| 444 | for i in reversed(axis): |
| 445 | x = f(builtin.Reduce(mode="sum", axis=i), x) |
| 446 | x = f(builtin.RemoveAxis(axis=[*reversed(axis)]), x) |
| 447 | return x |
| 448 | |
| 449 | ctx.dims2val = dims2val |
| 450 | ctx.reduce = reduce_sum |
| 451 | ctx.reshape = lambda x, dims: f(builtin.Reshape(), x, concat_dims(dims)) |
| 452 | ctx.matmul = lambda x, y: f(builtin.BatchedMatrixMul(), x, y) |
| 453 | ctx.broadcast = lambda x, dims: f(builtin.Broadcast(), x, concat_dims(dims)) |
| 454 | ctx.transpose = lambda x, axis: f(builtin.Dimshuffle(axis), x) |
| 455 | ctx.diag_plane = partial(diag_plane_subgraph, f=f, c=c) |
| 456 | return (einsum_impl(shapes, output_shape, inputs, ctx),), (True,) |
| 457 | |
| 458 | return einsum |
| 459 | |
| 460 | |
| 461 | def einsum_subgraph(equation, inputs: Tuple[Tensor, ...]) -> Tensor: |