int64_t NumElements(AbstractTensorHandle* tensor) const override;

  // Consumes references to the tensors in the gradient_tensors list and returns
  // a tensor with the result.
  AbstractTensorHandle* AggregateGradients(
      absl::Span<AbstractTensorHandle* const> gradient_tensors) const override;

  // Calls the passed-in backward function.
  // op_type is the op's name provided in RecordOperation.