absl::Span<AbstractTensorHandle*> grad_inputs) override {
    // Grad for Softmax Input
    TF_RETURN_IF_ERROR(BroadcastMul(
        ctx, grad_outputs[0], forward_outputs_[1],
        grad_inputs.subspan(0, 1)));  // upstream_grad * local softmax grad

    // Grad for labels is null
    grad_inputs[1] = nullptr;
    return absl::OkStatus();
  }

namespace gradients {
Status NotDifferentiableGradientFunction::Compute(
    AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> grad_outputs,
    absl::Span<AbstractTensorHandle*> grad_inputs) {
  for (int i = 0; i < grad_inputs.size(); i++) {
    grad_inputs[i] = nullptr;
  }
  return OkStatus();
}

Status RegisterNotDifferentiable(GradientRegistry* registry, const string& op) {

                 absl::Span<AbstractTensorHandle* const> grad_outputs,
                 absl::Span<AbstractTensorHandle*> grad_inputs) override {
    CHECK_EQ(grad_outputs.size(), 1);
    CHECK_EQ(grad_inputs.size(), 1);
    grad_inputs[0] = grad_outputs[0];
    if (grad_inputs[0]) {
      grad_inputs[0]->Ref();
    }
    return absl::OkStatus();
  }
};

// Computes:
//
// @tf.custom_gradient
// def f(input):

                 absl::Span<AbstractTensorHandle*> grad_inputs) override {
    // TODO(b/161805092): Support broadcasting.

    DCHECK(grad_outputs[0]);
    grad_inputs[0] = grad_outputs[0];
    grad_inputs[1] = grad_outputs[0];

    grad_inputs[0]->Ref();
    grad_inputs[1]->Ref();
    return absl::OkStatus();
  }
  ~AddGradientFunction() override {}
};

namespace tensorflow {
namespace gradients {
// Ignores `grad_outputs` and sets all entries in grad_inputs to nullptr.
class NotDifferentiableGradientFunction : public GradientFunction {
  Status Compute(AbstractContext* ctx,
                 absl::Span<AbstractTensorHandle* const> grad_outputs,
                 absl::Span<AbstractTensorHandle*> grad_inputs) override;
};
// Shorthand for registry->Register(op, new NotDifferentiableGradientFunction)

                 absl::Span<AbstractTensorHandle*> grad_inputs) override {
    for (int i = 0; i < grad_outputs.size(); i++) {
      auto grad_input = grad_outputs[i];
      // TODO(srbs): Should we add a copy contructor to AbstractTensorHandle
      // that takes care of this similar to `Tensor`?
      if (grad_input) {
        grad_input->Ref();
      }
      grad_inputs[i] = grad_input;
    }
    return absl::OkStatus();
  }

  }

  void AddGradients(bool grad_inputs_provided, const char* prefix,
                    TF_Output* inputs, int ninputs, TF_Output* outputs,
                    int noutputs, TF_Output* grad_outputs) {
    if (grad_inputs_provided) {
      TF_Output grad_inputs[1];
      const float grad_inputs_val[] = {1.0, 1.0, 1.0, 1.0};
      TF_Operation* grad_inputs_op =

//  public:
//   Status Compute(Context* ctx,
//                  absl::Span<AbstractTensorHandle* const> grad_inputs,
//                  absl::Span<AbstractTensorHandle*> grad_outputs) override {
//     grad_outputs[0] = grad_inputs[0];
//     grad_outputs[1] = grad_inputs[0];
//     grad_outputs[0]->Ref();
//     grad_outputs[1]->Ref();
//     return OkStatus();
//   }
//   ~AddGradientFunction() override {}