(equation, inputs)
| 371 | |
| 372 | |
| 373 | def einsum_interpret(equation, inputs): |
| 374 | input_ndims = tuple(map(lambda x: x.ndim, inputs)) |
| 375 | shapes, output_shape = einsum_parse_equation(equation, input_ndims) |
| 376 | ctx = EinsumContext() |
| 377 | dim2val_cache = {} |
| 378 | |
| 379 | def dim2val(dim: str): |
| 380 | if dim in dim2val_cache: |
| 381 | return dim2val_cache[dim] |
| 382 | for i, shape in enumerate(shapes): |
| 383 | if dim in shape: |
| 384 | dim2val_cache[dim] = inputs[i].shape[shape.index(dim)] |
| 385 | break |
| 386 | return dim2val_cache[dim] |
| 387 | |
| 388 | def dims2val(dims: str): |
| 389 | if len(dims) == 0: |
| 390 | return 1 |
| 391 | return reduce(lambda x, y: x * y, map(dim2val, dims)) |
| 392 | |
| 393 | ctx.dims2val = dims2val |
| 394 | ctx.reduce = lambda x, axis: sum(x, axis=axis) |
| 395 | ctx.reshape = reshape |
| 396 | ctx.matmul = matmul |
| 397 | ctx.broadcast = broadcast_to |
| 398 | ctx.transpose = transpose |
| 399 | ctx.diag_plane = diag_plane_interpret |
| 400 | return einsum_impl(shapes, output_shape, inputs, ctx) |
| 401 | |
| 402 | |
| 403 | @lru_cache(maxsize=None) |
no test coverage detected