Shared logic for transposed convolution emission. Handles weight transform, input/output transposition, bias, and node emission for all spatial dimensions. Called by both the specific conv_transpose handlers and the unified aten.convolution.default handler.
(
P: MLXProgramBuilder,
n: Node,
x_node: Node,
w_node: Node,
bias_node,
stride: list,
padding: list,
dilation: list,
output_padding: list,
groups: int,
ndim: int,
)
| 2459 | |
| 2460 | |
| 2461 | def _emit_conv_transpose( |
| 2462 | P: MLXProgramBuilder, |
| 2463 | n: Node, |
| 2464 | x_node: Node, |
| 2465 | w_node: Node, |
| 2466 | bias_node, |
| 2467 | stride: list, |
| 2468 | padding: list, |
| 2469 | dilation: list, |
| 2470 | output_padding: list, |
| 2471 | groups: int, |
| 2472 | ndim: int, |
| 2473 | ) -> Slot: |
| 2474 | """Shared logic for transposed convolution emission. |
| 2475 | |
| 2476 | Handles weight transform, input/output transposition, bias, and node emission |
| 2477 | for all spatial dimensions. Called by both the specific conv_transpose handlers |
| 2478 | and the unified aten.convolution.default handler. |
| 2479 | """ |
| 2480 | if ndim == 3 and groups != 1: |
| 2481 | raise ValueError( |
| 2482 | "conv_transpose with groups != 1 is not supported for 3D by MLX" |
| 2483 | ) |
| 2484 | |
| 2485 | w = _emit_conv_transpose_weight(P, w_node, groups, ndim=ndim) |
| 2486 | x, bias = P.slot_map([x_node, bias_node]) |
| 2487 | |
| 2488 | # Transpose input: channels-first -> channels-last |
| 2489 | ch_first_to_last = list(range(ndim + 2)) |
| 2490 | ch_first_to_last = [0] + list(range(2, ndim + 2)) + [1] |
| 2491 | ch_last_to_first = [0, ndim + 1] + list(range(1, ndim + 1)) |
| 2492 | |
| 2493 | _, tmp = P.make_tmp_slot() |
| 2494 | P.emit( |
| 2495 | TransposeNode(x=P.slot_to_tid(x), out=P.slot_to_tid(tmp), perm=ch_first_to_last) |
| 2496 | ) |
| 2497 | |
| 2498 | if ndim == 1: |
| 2499 | P.emit( |
| 2500 | ConvTranspose1DNode( |
| 2501 | x=P.slot_to_tid(tmp), |
| 2502 | w=P.slot_to_tid(w), |
| 2503 | out=P.slot_to_tid(tmp), |
| 2504 | stride=stride[0], |
| 2505 | padding=padding[0], |
| 2506 | dilation=dilation[0], |
| 2507 | output_padding=output_padding[0], |
| 2508 | groups=groups, |
| 2509 | ) |
| 2510 | ) |
| 2511 | elif ndim == 2: |
| 2512 | P.emit( |
| 2513 | ConvTranspose2DNode( |
| 2514 | x=P.slot_to_tid(tmp), |
| 2515 | w=P.slot_to_tid(w), |
| 2516 | out=P.slot_to_tid(tmp), |
| 2517 | stride_h=stride[0], |
| 2518 | stride_w=stride[1], |
no test coverage detected