def d_linear(input, weight, output_grad):
    """
    Inputs:
        - weight:      Dim (D2, D1)
        - bias:        Dim (D1,)
        - output_grad: Gradient of output w.r.t. to loss with dim (N, D2)
    Outputs:
        - dldx : Gradient w.r.t. input
        - dldw : Gradient w.r.t. weights
        - dldb : Gradient w.r.t. bias
    """
    dldx = output_grad.mm(weight)
    dldw = output_grad.t().mm(input)
    dldb = output_grad.sum(dim=0)
    return dldx, dldw, dldb