}
  }

  TF_DeleteTensor(analytical_tensor);
  delete[] danalytical;
}

Model BuildGradModel(Model forward, GradientRegistry registry) {
  return [forward_model = std::move(forward),
          grad_registry = std::move(registry)](
             AbstractContext* ctx,
             absl::Span<AbstractTensorHandle* const> inputs,
             absl::Span<AbstractTensorHandle*> outputs) -> Status {

class GraphOptPass
    : public mlir::PassWrapper<GraphOptPass,
                               mlir::OperationPass<mlir::ModuleOp>> {
  void getDependentDialects(mlir::DialectRegistry& registry) const override {
    mlir::RegisterAllTensorFlowDialects(registry);
  }

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GraphOptPass)

  explicit GraphOptPass(std::vector<tensorflow::GraphOptimizationPass*> passes)