// Handle the entry functions only.
  if (func.getName() != "main" && (!attrs || attrs.empty())) {
    return;
  }

  OpBuilder builder(func);
  FunctionType func_type = func.getFunctionType();
  llvm::SmallVector<Type, 4> new_input_types(func_type.getInputs().begin(),
                                             func_type.getInputs().end());

    func->removeAttr("tf_saved_model.exported_names");

    ArrayRef<Value> args;
    builder.create<func::CallOp>(session_initializer.getLoc(),
                                 func.getFunctionType().getResults(),
                                 func.getSymName(), args);
  }
  builder.create<func::ReturnOp>(session_initializer.getLoc());

  session_initializer.erase();
}

}  // namespace

void CreateRegionWithCall(func::FuncOp func, Region& region, Location loc) {
  OpBuilder builder(region);
  auto block = builder.createBlock(&region);
  SmallVector<Value, 4> new_operands;
  for (Type t : func.getFunctionType().getInputs())
    new_operands.push_back(block->addArgument(t, loc));
  auto call = builder.create<func::CallOp>(loc, func, new_operands);
  builder.create<YieldOp>(loc, call.getResults());

        returned_value.get().getDefiningOp()->emitError(
            "The producer of quantized type result should be a tfr.cast op.");
        return;
      }
    }

    auto new_func_type = builder.getFunctionType(block.getArgumentTypes(),
                                                 terminator->getOperandTypes());
    func.setType(new_func_type);
  });
}

// Creates an instance of the pass to decompose the TF ops.

      module_body.push_front(other);
      other.addEntryBlock();
      OpBuilder builder(other.getRegion());
      auto call_op = builder.create<mlir::func::CallOp>(
          f.getLoc(), f.getFunctionType().getResults(), f.getSymName(),
          other.getRegion().getArguments());
      builder.create<mlir::func::ReturnOp>(f.getLoc(), call_op.getResults());
    }

    Unexport(f);
  }
}

    // functions.
    auto func_names = mlir::tf_saved_model::GetExportedNames(func);
    if (func_names.empty()) return mlir::WalkResult::advance();

    auto func_type = func.getFunctionType();

    // Here we walk through each arguments and find out the input/output names,
    // and input devices, variables used by this function.
    llvm::SmallVector<llvm::StringRef, 4> input_names;
    llvm::SmallVector<

}

void LegalizeJaxRandomPass::runOnOperation() {
  auto func = getOperation();
  if (!IsJaxRandomUniform(func) && !IsJaxRandomNormal(func)) return;
  auto result_tuple_ty =
      mlir::dyn_cast_or_null<TupleType>(func.getFunctionType().getResult(0));
  if (!result_tuple_ty) return;
  if (result_tuple_ty.size() != 1) return;
  auto result_ty = mlir::dyn_cast<ShapedType>(result_tuple_ty.getType(0));

  func.eraseBody();
  func.addEntryBlock();

func::FuncOp createMaxUnpoolingFunc(
    mlir::Builder* builder, const SmallVector<mlir::Type, NInput>& input_types,
    const SmallVector<mlir::Type, NOutput>& output_types) {
  auto func_type = builder->getFunctionType(input_types, output_types);
  auto func = func::FuncOp::create(
      mlir::NameLoc::get(builder->getStringAttr("fused_func")), "fused_func",
      func_type, {});

  func.addEntryBlock();

void UpdateFuncType(func::FuncOp func) {
  Operation* terminator = func.front().getTerminator();
  auto return_types = llvm::to_vector<4>(terminator->getOperandTypes());

  FunctionType func_type = func.getFunctionType();
  if (llvm::ArrayRef(return_types) == func_type.getResults()) return;

  auto updated_type =
      FunctionType::get(func.getContext(), func_type.getInputs(), return_types);
  func.setType(updated_type);
}

    def __init__(self, *args, **kwargs) -> None: ...

class Block_Iterator:
    def __init__(self, *args, **kwargs) -> None: ...

class Builder:
    def __init__(self, arg0: MLIRContext) -> None: ...
    def getFunctionType(self, arg0: list[Type], arg1: list[Type]) -> FunctionType: ...

class FloatType(Type):
    def __init__(self, *args, **kwargs) -> None: ...
    def getBF16(self) -> FloatType: ...
    def getF16(self) -> FloatType: ...