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

    // Grad for bias
    std::string name = "bias_add_grad";
    TF_RETURN_IF_ERROR(BiasAddGrad(ctx, upstream_grad, &grad_inputs[1],
                                   data_format.c_str(), name.c_str()));

    return absl::OkStatus();
  }
  ~BiasAddGradientFunction() override {}