// Start step 0
C_0 = cond(args_0)
N_0 = non_tpu(args_0)
if (C_0) {
   F_0 = forward(args_0, N_0)
   T_0 = core_tpu(args_0, N_0, F_0)
   // B_0 = backward() is not evaluated here.
}

args_1 = update_args(args_0, N_0, T_0)

// Start step 1
C_1 = cond(args_1)
N_1 = non_tpu(args_1)
if (C_1) {
   F_1 = forward(args_1, N_1)
   // T_1 = core_tpu() is not evaluated here.

// CHECK:             %{{.*}}, %[[STACK_PUSH_CONTROL:.*]] = tf_executor.island wraps "tf.StackPushV2"(%[[STACK_HANDLE]], %[[ARG_1]]) : (tensor<!tf_type.resource<tensor<f32>>>, tensor<f32>) -> tensor<f32>
// CHECK:             %[[ADD:.*]], %{{.*}} = tf_executor.island wraps "tf.Add"(%[[ARG_1]], %[[ARG_1]]) : (tensor<f32>, tensor<f32>) -> tensor<f32>

// CHECK-SAME: %[[ARG_1]] as %{{[a-z0-9]*}}
// CHECK-SAME: %[[ARG_3]] as %{{[a-z0-9]*}}
// CHECK-SAME: %[[ARG_5]] as %{{[a-z0-9]*}}


// Test that the `is_mirrored_variable` attribute is preserved in the
// tf_device.replicate op.
// CHECK-LABEL: func @mirrored_variables

    func.return
 }
 ```

 with the following reduction function:

 ```mlir
 func.func private @__reduce_func_1(%arg0: tensor<i64> {tf._user_specified_name = "args_0"},
   %arg1: tensor<32xf32> {tf._user_specified_name = "args_1"}) -> (tensor<i64>)
   attributes {tf._tf_data_function = true, tf.signature.is_stateful} {
     %0 = "tf.JustPretend"(%arg1) : (tensor<32xf32>) -> (tensor<i64>)
     func.return %0 : tensor<i64>