%cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> {device = ""} : () -> tensor<i32>
  %0 = "tf.MultiDeviceIteratorFromStringHandle"(%arg0) <{output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]}> {device = ""} : (tensor<!tf_type.string>) -> tensor<!tf_type.resource>

  ElementsAttr attr;
  if (auto cst = dyn_cast<mlir::arith::ConstantOp>(inst)) {
    // arith::ConstantOp have ElementAttr at this point due to validation of the
    // TFLite module.
    attr = mlir::cast<ElementsAttr>(cst.getValue());
  } else if (auto cst = dyn_cast<mlir::TF::ConstOp>(inst)) {
    attr = cst.getValue();
  } else if (auto cst = dyn_cast<tfl::ConstOp>(inst)) {
    attr = cst.getValue();

    module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
      func.func @main() -> (tensor<2xi32>) {
        %cst = "tf.Const"() {value = dense<[524170, 523952]> : tensor<2xi32>} : () -> tensor<2xi32>
        return %cst : tensor<2xi32>
    }
  })";

  auto compilation_result = CompileMlirModule(
      kHasReturnValuesAndNoMetadataRetvals,

      "tf.MakeIterator"(%arg0, %0) : (tensor<!tf_type.variant>, tensor<!tf_type.resource>) -> ()
      %cst = "tf.Const"() {value = dense<true> : tensor<i1>} : () -> tensor<i1>
      %1:2 = "tf.WhileRegion"(%cst, %arg1) ({
      ^bb0(%arg2: tensor<i1>, %arg3: tensor<i64>):
        "tf.Yield"(%arg2) : (tensor<i1>) -> ()
      }, {
      ^bb0(%arg2: tensor<i1>, %arg3: tensor<i64>):

    auto cst =
        rewriter.create<mhlo::ConstantOp>(op.getLoc(), zero_attr).getResult();
    auto split_updates = rewriter.create<TF::SplitOp>(
        op.getLoc(), split_updates_type, cst, updates);

    SmallVector<Value, 6> input_indices;
    input_indices.resize(input_type.getRank(), cst);

    for (auto pair :