// `tf.AssignVariableOp`s and `tf.Const`s. Collects and returns the
// `tf.VarHandleOp`s that are initialized by these `tf.AssignVariableOp`s.
SmallVector<TF::VarHandleOp> CollectVariableOps(
    func::FuncOp session_init_func) {
  SmallVector<TF::VarHandleOp> var_handle_ops{};

  for (auto assign_variable_op : llvm::make_early_inc_range(
           session_init_func.getOps<TF::AssignVariableOp>())) {

  Value output_layer_norm_coefficients_;

  mlir::TFL::LSTMOp lstm_;

  Value none_;
  SmallVector<int64_t, 1> bias_slice_shape_;
  SmallVector<int64_t, 1> bias_size_values_;
  SmallVector<int64_t, 2> weight_slice_shape_;
  SmallVector<int64_t, 2> weight_slice_size_input_values_;
  SmallVector<int64_t, 2> weight_slice_size_recurrent_values_;
  OpBuilder builder_;
};

      : input_shapes_(input_shapes) {}
  void runOnOperation() override {
    // Parse `input_arg_shapes_` if provided (test only)
    SmallVector<ArrayRef<int64_t>> input_shapes_vec;
    absl::StatusOr<SmallVector<SmallVector<int64_t>>> parsed_shapes;
    if (!input_arg_shapes_.empty()) {
      parsed_shapes = ParseArgumentShapes(input_arg_shapes_);
      if (!parsed_shapes.ok()) {

    mlir::OpBuilder* builder,
    llvm::SmallVectorImpl<llvm::SmallVector<mlir::Value, 4>>* input_list);

// Extracts a list of OpSharding that represent output sharding configuration of
// `tf_device.cluster`.
mlir::LogicalResult ParseAndValidateOutputSharding(
    int num_cores_per_replica, mlir::tf_device::ClusterFuncOp cluster_func,
    mlir::SmallVector<xla::OpSharding, 4>* output_sharding_list);

#include <type_traits>

#include "llvm/ADT/ArrayRef.h"  // IWYU pragma: keep; required to include the definition of ArrayRef
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"  // IWYU pragma: keep; required to include the definition of SmallVector
#include "mlir/Support/LLVM.h"  // from @llvm-project

namespace mlir::quant {

// Permutes `values` with `permutation`. Returns the permuted values. Sizes of

  SmallVector<float> TransposeValues(const ArrayRef<float> values) const {
    SmallVector<float> transposed_values(values.size());
    SmallVector<int64_t> current_indices = {};
    TransposeRecursively(values, transposed_values, current_indices);

    return transposed_values;
  }

  // Returns the shape after permutation.
  SmallVector<int64_t> GetTargetShape() const { return target_shape_; }

  void runOnOperation() override;
};

LogicalResult HasNoNestedEntryFunctions(
    const llvm::SmallVector<func::FuncOp> &entry_funcs, SymbolTable &symtab) {
  auto calls_entry_functions = [&](SymbolUserOpInterface op) {
    llvm::SmallVector<func::FuncOp> callees;
    if (GetCallees(op, symtab, callees).failed()) {
      return false;
    }
    for (auto &callee : callees) {

// BlockArg or as a result of an Operation.
llvm::SmallVector<Value> AccumulateOperandsDefinedAbove(
    const llvm::SetVector<Operation*>& partition_ops);

// Similar to `AccumulateOperandsDefinedAbove()`, computes the Value(s) that are
// defined within a Subgraph and referenced in a descendant Operation. These
// Values(s) are to be returned by the new raised function.
llvm::SmallVector<Value> AccumulateResultsDefinedWithin(

  bool rewrite = false;
  const int num_regions = parallel_execute.getNumRegions();
  llvm::SmallVector<Value, 4> results_to_remap;

  // Go through each region and find AssignVariableOps that can be moved into
  // the parallel_execute region. Result indices by region index are collected,
  // so they can be removed afterwards.
  llvm::SmallVector<llvm::SmallVector<int, 4>, 4> results_to_remove_by_region;
  results_to_remove_by_region.resize(num_regions);

bool IsEntryFunction(func::FuncOp func);

// Get all the entry functions in an MLIR module.
llvm::SmallVector<func::FuncOp> GetEntryFunctions(ModuleOp module);

// Get all the functions referenced in a symber user op and save them in
// `callees`.
LogicalResult GetCallees(SymbolUserOpInterface op, SymbolTable &symtab,
                         llvm::SmallVector<func::FuncOp> &callees);