@triton.jit
def _mlp_wide_kernel_bwd_dx(pid_h, pid_b, x_ptr, w1_ptr, w2_ptr, o_ptr,
    dx_ptr, dw1_ptr, dw2_ptr, do_ptr, H, B, D: tl.constexpr, E, stride_xb,
    stride_xd, stride_w1d, stride_w1e, stride_w2e, stride_w2d, stride_ob,
    stride_od, stride_dxb, stride_dxd, stride_dw1d, stride_dw1e,
    stride_dw2e, stride_dw2d, stride_dob, stride_dod, BLOCK_SIZE_B: tl.
    constexpr, BLOCK_SIZE_E: tl.constexpr, ACTIVATION: tl.constexpr):
    """Kernel for computing the mlp_bwd_dx
    Z = X @ W1, H = f(Z), O = H @ W2
    - X has shape (B, D)
    - W1 has shape (D, E)
    - W2 has shape (E, D)
    - O has shape (B, D)
    - dX has shape (B, D)
    - dW1 has shape (D, E)
    - dW2 has shape (E, D)
    - dO has shape (B, D)
    """
    TARGET_TYPE = x_ptr.type.element_ty
    offs_b = tl.arange(0, BLOCK_SIZE_B)
    offs_d = tl.arange(0, D)
    offs_e = tl.arange(0, BLOCK_SIZE_E)
    x_ptrs = x_ptr + ((pid_h * B + pid_b * BLOCK_SIZE_B + offs_b[:, None]) *
        stride_xb + offs_d[None, :] * stride_xd)
    x_mask = (offs_b[:, None] < B - pid_b * BLOCK_SIZE_B) & (offs_d[None, :
        ] < D)
    do_ptrs = do_ptr + ((pid_h * B + pid_b * BLOCK_SIZE_B + offs_b[:, None]
        ) * stride_dob + offs_d[None, :] * stride_dod)
    do_mask = (offs_b[:, None] < B - pid_b * BLOCK_SIZE_B) & (offs_d[None,
        :] < D)
    w1_ptrs = w1_ptr + ((pid_h * D + offs_d[:, None]) * stride_w1d + offs_e
        [None, :] * stride_w1e)
    w2_ptrs = w2_ptr + ((pid_h * E + offs_e[:, None]) * stride_w2e + offs_d
        [None, :] * stride_w2d)
    dw1_ptrs = dw1_ptr + ((pid_h * D + offs_d[:, None]) * stride_dw1d + 
        offs_e[None, :] * stride_dw1e)
    dw2_ptrs = dw2_ptr + ((pid_h * E + offs_e[:, None]) * stride_dw2e + 
        offs_d[None, :] * stride_dw2d)
    x = tl.load(x_ptrs, mask=x_mask, other=0.0)
    do = tl.load(do_ptrs, mask=do_mask, other=0.0)
    dx = tl.zeros((BLOCK_SIZE_B, D), dtype=tl.float32)
    for e in range(0, tl.cdiv(E, BLOCK_SIZE_E)):
        w1_mask = (offs_d[:, None] < D) & (offs_e[None, :] < E - e *
            BLOCK_SIZE_E)
        w2_mask = (offs_e[:, None] < E - e * BLOCK_SIZE_E) & (offs_d[None,
            :] < D)
        w1 = tl.load(w1_ptrs, mask=w1_mask, other=0.0)
        w2 = tl.load(w2_ptrs, mask=w2_mask, other=0.0)
        z = tl.dot(x, w1, out_dtype=tl.float32)
        if ACTIVATION == 'leaky_relu':
            h = leaky_relu(z).to(TARGET_TYPE)
        else:
            h = z.to(TARGET_TYPE)
        dh = tl.dot(do, tl.trans(w2), out_dtype=tl.float32)
        if ACTIVATION == 'leaky_relu':
            dz = (dh * d_leacky_relu_inv_backward(z)).to(TARGET_TYPE)
        else:
            dz = dh.to(TARGET_TYPE)
        dx += tl.dot(dz, tl.trans(w1), out_dtype=tl.float32)
        w1_ptrs += BLOCK_SIZE_E * stride_w1e
        w2_ptrs += BLOCK_SIZE_E * stride_w2e
        dw1_ptrs += BLOCK_SIZE_E * stride_dw1e
        dw2_ptrs += BLOCK_SIZE_E * stride_dw2e
    return dx
