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();
  }

                 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 absl::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):

Status AddWhileGradientLoop(WhileContext* while_ctx,
                            const std::vector<Output>& grad_inputs,
                            const Output& backprop_execution_pred,
                            const Scope& parent_scope,
                            std::vector<Output>* grad_outputs) {
  DCHECK_EQ(grad_inputs.size(), while_ctx->body_outputs().size());
  DCHECK_EQ(while_ctx->body_inputs().size(), while_ctx->body_outputs().size());

                        const std::vector<Output>& grad_inputs,
                        std::vector<Output>* grad_outputs) {
  ops::GradFunc grad_fn;
  TF_RETURN_IF_ERROR(ops::GradOpRegistry::Global()->Lookup(
      op.node()->type_string(), &grad_fn));
  TF_RETURN_IF_ERROR(grad_fn(scope, op, grad_inputs, grad_outputs));
  TF_RETURN_IF_ERROR(scope.status());
  return absl::OkStatus();
}