Function _emit_conv_bias

backends/mlx/ops.py:2287–2313

Reshape conv bias to channel-last broadcast shape and add to tmp in-place. After the convolution the activation is in channel-last layout, so the bias (shape ``[C_out]``) must be reshaped to ``[1, …, 1, -1]`` with *ndim* dimensions before being added. Does nothing when *bias* is ``None

(
    P: MLXProgramBuilder, bias: Optional[Slot], tmp: Slot, ndim: int
)

def _emit_conv_bias(
    P: MLXProgramBuilder, bias: Optional[Slot], tmp: Slot, ndim: int
) -> None:
    """Reshape conv bias to channel-last broadcast shape and add to tmp in-place.

    After the convolution the activation is in channel-last layout, so the bias
    (shape ``[C_out]``) must be reshaped to ``[1, …, 1, -1]`` with *ndim*
    dimensions before being added.  Does nothing when *bias* is ``None``.
    """
    if bias is None:
        return
    _, tmp2 = P.make_tmp_slot()
    shape = [IntOrVid.from_literal(1)] * (ndim - 1) + [IntOrVid.from_literal(-1)]
    P.emit(
        ReshapeNode(
            x=P.slot_to_tid(bias),
            out=P.slot_to_tid(tmp2),
            shape=shape,
        )
    )
    P.emit(
        AddNode(
            a=P.slot_to_tid(tmp),
            b=P.slot_to_tid(tmp2),
            out=P.slot_to_tid(tmp),
        )
    )


def _emit_conv(