}
    }

    // Special type fixes for scalar tensor types.
    // TFLite flatbuffer schema doesn't distinguish scalar tensor shapes
    // and unranked tensor shapes (i.e. they are both represented as an empty
    // INT32 list), see b/138865275. MLIR importer conservatively treats them as
    // unranked tensor types. Here we set them to scalar tensor types when it is
    // safe.

  }
  return %graph#0, %graph#1, %graph#2, %graph#3 : !tf_res, !tf_res, tensor<f32>, tensor<f32>
}
```

After:

```mlir
func @while_body(%arg0: !tf_res, %arg1: !tf_res, %arg2: tensor<f32>, %arg3: tensor<f32>, %chain_0: tensor<i32>, %chain_1: tensor<i32>) -> (!tf_res, !tf_res, tensor<f32>, tensor<f32>, tensor<i32>, tensor<i32>) {
  %graph:6 = tf_executor.graph {

}

# MLIR-LABEL: func @main
# MLIR-SAME:  (%[[ARG_0:[a-z0-9]+]]: tensor<1x1x1x256x!quant.uniform<i8:f32, 0.21632751372549019:27>>) -> tensor<1x6x31x!quant.uniform<i8:f32, 0.09363494573854933:22>>
# MLIR-SAME:  control_outputs = ""
# MLIR-SAME:  inputs = "input"
# MLIR-SAME:  outputs = "output"
# MLIR:         %[[shape:.*]] = arith.constant dense<[1, -1, 31]> : tensor<3xi32>

// quantized types.
std::unique_ptr<OperationPass<ModuleOp>>
CreateComposeUniformQuantizedTypePass();

// Creates a pass that finds stablehlo ops that accept or produce uniform
// quantized typed tensors and converts them to equivalent ops in the TFLite
// dialect.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateUniformQuantizedStableHloToTflPass();

// Create a pass that commute transposes through specific ops

  TFE_DeleteContextOptions(opts);

  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());

  /* Want to test simple MatMul example with abstract tensors:
    [[1,2],   *  [[5,6],   =   [[19,22],
     [3,4]]       [7,8]]        [43,50]]
  */

  // Build 1st Matrix.
  int64_t dims[] = {2, 2};  // Matrices will be 2 x 2

versions {
  producer: 27
}

# CHECK:      {
# CHECK-NEXT:   version: 3,
# CHECK-NEXT:   operator_codes: [ {
# CHECK-NEXT:     version: 1
# CHECK-NEXT:   } ],
# CHECK-NEXT:   subgraphs: [ {
# CHECK-NEXT:     tensors: [ {
# CHECK-NEXT:       shape: [ 4 ],
# CHECK-NEXT:       type: INT32,
# CHECK-NEXT:       buffer: 1,
# CHECK-NEXT:       name: "input0",
# CHECK-NEXT:       quantization: {
# CHECK-EMPTY:
# CHECK-NEXT:       },

static opt<std::string> input_arrays_flag(
    "input-arrays",
    llvm::cl::desc(
        "List of input tensors, if different from the default inputs"),
    llvm::cl::init(""));

// NOLINTNEXTLINE
static opt<std::string> output_arrays_flag(
    "output-arrays",
    llvm::cl::desc(
        "List of output tensors, if different from the default outputs"),
    llvm::cl::init(""));

using llvm::cl::opt;

  return absl::OkStatus();
}

Tensor CreateStringTensor(const string& value) {
  Tensor tensor(DT_STRING, TensorShape({}));
  tensor.scalar<tstring>()() = value;
  return tensor;
}

void AddAssetsTensorsToInputs(const StringPiece export_dir,
                              const std::vector<AssetFileDef>& asset_file_defs,
                              std::vector<std::pair<string, Tensor>>* inputs) {

                                  llvm::cl::init("-"));
// NOLINTNEXTLINE
opt<bool> use_splatted_constant(
    "use-splatted-constant",
    llvm::cl::desc(
        "Replace constants with randomly generated splatted tensors"),
    llvm::cl::init(false), llvm::cl::Hidden);
// NOLINTNEXTLINE
opt<bool> input_mlir(
    "input-mlir",
    llvm::cl::desc("Take input TensorFlow model in textual MLIR instead of "

    return {0, 2, 3, 1};
  } else {
    return {};
  }
}

// Shuffle elements in the `attr` according to the permutation. Optional
// `inner_size` allows to shuffle array attributes created from rank 2 tensors
// on outer dimension only.
ArrayAttr ShuffleArrayAttr(ArrayAttr attr, ArrayRef<int64_t> permutation,
                           int inner_size) {
  if (attr.empty()) return attr;

  assert(attr.size() % inner_size == 0);