ctx_.reset(ctx_raw);
  }
};

// TODO(b/191361582): Use Abstract* APIs for building functions so that we can
// test with MLIR.
FunctionDef SquareFunc() {
  return FunctionDefHelper::Define(
      // Function Name
      "SquareFunc",
      // Args
      {"x: float"},
      // Returns
      {"y: float"},
      // Attr def
      {},
      // Nodes
      {{/*ret=*/{"y"},