%0 = arith.constant dense<0> : tensor<3xi32>
  %1 = "tfl.dynamic_update_slice"(%arg0, %arg1, %0) : (tensor<2x7x14xf32>, tensor<2x7x7xf32>, tensor<3xi32>) -> tensor<2x7x14xf32>
  // CHECK: "tfl.dynamic_update_slice"
  func.return %1 : tensor<2x7x14xf32>
}

// CHECK-LABEL: OptimizeTranposeWithRank7orMoreEffectiveRank6

    func.return %1 : tensor<2x2xf32>
  }

  // CHECK-LABEL: dynamic_update_slice
  // CHECK-SAME: (%[[ARG0:.*]]: tensor<3x4xi32>, %[[ARG1:.*]]: tensor<2x2xi32>, %[[ARG2:.*]]: tensor<2xi32>
  func.func @dynamic_update_slice(%arg0: tensor<3x4xi32>, %arg1: tensor<2x2xi32>, %arg2: tensor<2xi32>) -> tensor<3x4xi32> {

    // CHECK: %[[SLICE0:.*]] = "mhlo.slice"(%[[ARG2]])

//CHECK-NEXT:}

func.func @dynamic_update_slice(%arg0: tensor<4x4xi64>, %arg1: tensor<2x3xi64>, %arg2: tensor<i64>, %arg3: tensor<i64>) -> tensor<4x4xi64> {
  %0 = "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (tensor<4x4xi64>, tensor<2x3xi64>, tensor<i64>, tensor<i64>) -> tensor<4x4xi64>
  return %0 : tensor<4x4xi64>
}

  READ_VARIABLE = 143,
  ASSIGN_VARIABLE = 144,
  BROADCAST_ARGS = 145,
  RANDOM_STANDARD_NORMAL = 146,
  BUCKETIZE = 147,
  RANDOM_UNIFORM = 148,
  MULTINOMIAL = 149,
  GELU = 150,
  DYNAMIC_UPDATE_SLICE = 151,
  RELU_0_TO_1 = 152,
  UNSORTED_SEGMENT_PROD = 153,
  UNSORTED_SEGMENT_MAX = 154,
  UNSORTED_SEGMENT_SUM = 155,
  ATAN2 = 156,
  UNSORTED_SEGMENT_MIN = 157,
  SIGN = 158,
  BITCAST = 159,

  // Assert that the tensor list operation is lowered to something that doesn't
  // get stuck on a broken dynamic update slice.
  ASSERT_THAT(compilation_result,
              Not(ComputationProtoContains("%.*=.*DynamicUpdateSlice")));
}

TEST(LegalizeTFTest, LegalizesFunctionWithBoundedDynamicArg) {
  static constexpr char kMlirModuleWithBoundedDynamicArgStr[] = R"(

        "tensorlist op lowering (default true)">,
      Option<"enable_dynamic_update_slice_", "enable-dynamic-update-slice",
             "bool", "false", "When specified to true, lower TensorListSetItem with "
             "DynamicUpdateSlice op (default false).">,
  ];
}

def ModifyIONodesPass : Pass<"tfl-modify-io-nodes", "mlir::func::FuncOp"> {
  let summary = "Modify the type of the model io nodes";
  let description = [{