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();
  }
  ~SparseSoftmaxCrossEntropyWithLogitsGradientFunction() override {}