*    dA =  U
     *    dB = -U
     *
     */

    // Grad for A
    DCHECK(grad_outputs[0]);
    grad_inputs[0] = grad_outputs[0];
    grad_inputs[0]->Ref();

    // Grad for B
    // negate the upstream grad
    std::string name = "Neg_Sub_Grad_B";
    TF_RETURN_IF_ERROR(
        ops::Neg(ctx, grad_outputs[0], &grad_inputs[1], name.c_str()));

    return absl::OkStatus();
  }