namespace gradients {
namespace {

class ReluGradientFunction : public GradientFunction {
 public:
  explicit ReluGradientFunction(vector<AbstractTensorHandle*> f_outputs)
      : forward_outputs_(f_outputs) {
    for (auto output : forward_outputs_) {
      if (output) {
        output->Ref();
      }
    }
  }

  Status Compute(AbstractContext* ctx,
                 absl::Span<AbstractTensorHandle* const> grad_outputs,

                                    vector<AbstractTensorHandle*> f_outputs)
      : forward_inputs_(f_inputs), forward_outputs_(f_outputs) {
    for (auto input : forward_inputs_) {
      if (input) {
        input->Ref();
      }
    }
    for (auto output : forward_outputs_) {
      if (output) {
        output->Ref();
      }
    }
  }

  Status Compute(AbstractContext* ctx,