(
inp: Tensor, inp_ndim: int, axes: List[int], f: Any, c: Any
)
| 348 | |
| 349 | |
| 350 | def diag_plane_subgraph( |
| 351 | inp: Tensor, inp_ndim: int, axes: List[int], f: Any, c: Any |
| 352 | ) -> Any: |
| 353 | from megengine.core.ops import builtin |
| 354 | |
| 355 | all_axes = set(range(inp_ndim)) |
| 356 | remaining_axes = sorted(all_axes.difference(axes)) |
| 357 | transposed = f(builtin.Dimshuffle(list(axes) + remaining_axes), inp) |
| 358 | diag_len = f(builtin.GetVarShape(axis=axes[0]), inp) |
| 359 | mav_arg = list(map(lambda x: (x, False, False, False, True), range(len(axes)))) |
| 360 | mav_index = lambda: f( |
| 361 | builtin.TypeCvt("int32"), |
| 362 | f( |
| 363 | builtin.Linspace(), |
| 364 | c(0), |
| 365 | f(builtin.Elemwise("SUB"), diag_len, c(1)), |
| 366 | diag_len, |
| 367 | ), |
| 368 | ) |
| 369 | mav_indices = [mav_index() for _ in range(len(axes))] |
| 370 | return f(builtin.IndexingMultiAxisVec(mav_arg), transposed, *mav_indices) |
| 371 | |
| 372 | |
| 373 | def einsum_interpret(equation, inputs): |
nothing calls this directly
no test coverage detected