"tf.AssignAddVariableOp"(%37, %cst_12) {_has_manual_control_dependencies = true, _replication_info = "while/cluster_while_body_451", _xla_compile_device_type = "TPU", device = ""} : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<i64>) -> ()

    return %res_n, %arg1 : tensor<i32>, tensor<*x!tf_type.resource<tensor<i64>>>
  }

// CHECK: %[[ADD:.+]] = stablehlo.add %[[CONV]], %[[BIAS_CONST]] : tensor<1x5x5x4xf32>
// CHECK: %[[REDUCE_WINDOW_MAX:.+]] = "stablehlo.reduce_window"(%[[ADD]], %[[INIT_VALUE_CONST:.+]])
// CHECK: <{window_dimensions = array<i64: 1, 2, 2, 1>, window_strides = array<i64: 1, 2, 2, 1>}>
// CHECK: stablehlo.maximum

  "tf.InitializeTableFromTextFileV2"(%0, %cst) {delimiter = " ", device = "", key_index = -2 : i64, value_index = -1 : i64, vocab_size = -1 : i64} : (tensor<!tf_type.resource>, tensor<!tf_type.string>) -> ()
  func.return
  // CHECK-LABEL: func @init_all_tables

{_xla_inferred_shapes = [#tf_type.shape<*>, #tf_type.shape<*>], allowed_batch_sizes = [64, 128, 256], batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", enable_large_batch_splitting = true, f = @compute_0, max_batch_size = 256 : i64, max_enqueued_batches = 10000 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "computation"} : (tensor<?x?xi32>, tensor<?x?xi32>) -> (tensor<*xi32>, tensor<*xi32>)...

  // CHECK-SAME: begin_mask = 7 : i64
  // CHECK-SAME: ellipsis_mask = 0 : i64
  // CHECK-SAME: end_mask = 14 : i64
  // CHECK-SAME: new_axis_mask = 0 : i64
  // CHECK-SAME: shrink_axis_mask = 0 : i64

        quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
      } : (tensor<?xf32>, tensor<f32>, tensor<i32>) -> tensor<?x!tf_type.qint8>
      %1 = "tf.UniformDequantize"(%0, %scales, %zps) {
        quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
      } : (tensor<?x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<?xf32>

            tf_device.return %5: tensor<i64>
          }) {device = "/TPU:2"} : () -> (tensor<i64>)
          tf_device.return %4 : tensor<i64>
        }) : () -> (tensor<i32>, tensor<i64>)
        tf_device.return %3#0, %3#1 : tensor<i32>, tensor<i64>
      }
      tf_executor.yield %2#0, %2#1, %2#2, %2#3 : tensor<i32>, tensor<i32>, tensor<i64>, tensor<i64>
    }

  "tf.InitializeTableFromTextFileV2"(%0, %cst) {delimiter = " ", device = "", key_index = -2 : i64, value_index = -1 : i64, vocab_size = -1 : i64} : (tensor<!tf_type.resource>, tensor<!tf_type.string>) -> ()
  func.return
  // CHECK-LABEL: func @init_all_tables

  // CHECK-SAME: allowed_batch_sizes = [6]
  // CHECK-SAME: batch_timeout_micros = 100000 : i64
  // CHECK-SAME: batching_queue = ""
  // CHECK-SAME: container = ""
  // CHECK-SAME: enable_large_batch_splitting = false
  // CHECK-SAME: max_batch_size = 6 : i64
  // CHECK-SAME: max_enqueued_batches = 10 : i64
  // CHECK-SAME: num_batch_threads = 1 : i64
  // CHECK-SAME: shared_name = "batch/"

    tf_device.return %arg0, %arg1, %arg2, %arg3, %arg4 : tensor<i1>, tensor<i32>, tensor<i64>, tensor<f32>, tensor<f64>
  // CHECK-NEXT: }, {
  }, {
    // CHECK-NEXT: tf_device.return
    tf_device.return
  // CHECK-NEXT: }) : () -> (tensor<i1>, tensor<i64>, tensor<f64>)
  }) : () -> (tensor<i1>, tensor<i32>, tensor<i64>, tensor<f32>, tensor<f64>)