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

                                   std::vector<Output>* grad_outputs) {
  string kernel_type;
  TF_RETURN_IF_ERROR(
      GetNodeAttr(op.node()->attrs(), "kernel_type", &kernel_type));
  bool antialias;
  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "antialias", &antialias));
  grad_outputs->push_back(internal::ScaleAndTranslateGrad(
      scope, grad_inputs[0], op.input(0), op.input(2), op.input(3),

/// Implementations should add operations to compute the gradient outputs of
/// 'op' (returned in 'grad_outputs') using 'scope' and 'grad_inputs'.
typedef Status (*GradFunc)(const Scope& scope, const Operation& op,
                           const std::vector<Output>& grad_inputs,
                           std::vector<Output>* grad_outputs);

/// GradOpRegistry maintains a static registry of gradient functions.

                 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) {

  const int num_inputs = op.num_inputs();
  func_inputs.reserve(num_inputs + grad_inputs.size());
  input_dtypes.reserve(num_inputs);

  for (int i = 0; i < num_inputs; i++) {
    func_inputs.push_back(op.input(i));
    input_dtypes.push_back(op.input_type(i));
  }

  func_inputs.insert(std::end(func_inputs), std::begin(grad_inputs),
                     std::end(grad_inputs));

  auto grad = SymbolicGradient(scope, func_inputs, input_dtypes, f);

namespace tensorflow {

// Adds the gradient computation for the while loop associated with
// `while_ctx`. `grad_inputs` are the partial derivatives w.r.t. the loop
// outputs, i.e. the exit nodes.  The partial derivatives w.r.t. the loop
// inputs, i.e. the input loop vars, are returned in `grad_outputs`.
// `grad_inputs` and `grad_outputs` are both in loop-variable order, as defined
// by the original inputs to BuildWhileLoop().

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

                            const std::vector<Output>& inputs,
                            const std::vector<Output>& grad_inputs,
                            std::vector<Output>* grad_outputs);

// Same as above, but uses 'OnesLike' for all shapes in
// 'outputs' as grad_inputs.
Status AddSymbolicGradients(const Scope& scope,
                            const std::vector<Output>& outputs,

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