%3 = "tf.XlaCallModule"(%2) <{Sout = [#tf_type.shape<6x5>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @main_1, _stablehlo_module_attrs = {}, device = ""} : (tensor<3x10xf32>)...

  %conv = stablehlo.convolution(%quant_arg, %quant_cst) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3x!quant.uniform<i8:f32, 0.0039207626791561354:-128>>, tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.015>>) -> tensor<1x3x4x2x!quant.uniform<i32:f32, 5.8949912267181218E-5>>

    %2 = "tf.OpC"(%arg3) {is_stateless = true} : (tensor<*xf32>) -> tensor<*xf32>
    "tf.Yield"(%1, %2) : (tensor<*xf32>, tensor<*xf32>) -> ()
  }) {is_stateless = true, parallel_iterations = 10 : i64} : (tensor<*xf32>, tensor<*xf32>) -> (tensor<*xf32>, tensor<*xf32>)
  return %0#1 : tensor<*xf32>
}

// -----

// Check that we remove the second result and operand from the loop because the

          "Failed at tf.StridedSlice op\nsee current operation: %2 = "
          "\"tf.StridedSlice\"(%arg0, %1, %1, %0) <{begin_mask = 11 : "
          "i64, ellipsis_mask = 0 : i64, end_mask = 11 : i64, new_axis_mask = "
          "4 : i64, shrink_axis_mask = 0 : i64}> {device = \"\"} : "
          "(tensor<*xf32>, tensor<4xi32>, tensor<4xi32>, tensor<4xi32>) "
          "-> tensor<*xf32>\nError code: ERROR_NEEDS_FLEX_OPS",

    %0:3 = "tf.Unpack"(%arg0) {axis = 0 : i64} : (tensor<3xi64>) -> (tensor<i64>, tensor<i64>, tensor<i64>)
    return
  }
})";

  %0 = "tf.FakeQuantWithMinMaxArgs"(%arg1) {device = "", max = 2.000000e+00 : f32, min = -1.000000e+00 : f32, narrow_range = false, num_bits = 8 : i64} : (tensor<2x3x3x2xf32>) -> tensor<*xf32>
  %1 = "tf.FakeQuantWithMinMaxArgs"(%arg0) {device = "", max = 2.000000e-01 : f32, min = -1.000000e-01 : f32, narrow_range = false, num_bits = 8 : i64} : (tensor<1x3x4x3xf32>) -> tensor<*xf32>

  // CHECK-SAME: quant.uniform<i4:f32, 0.13329366842905679:-1>
  // CHECK: %[[dq2:.*]] = "quantfork.dcast"(%[[q2]])
  // CHECK-SAME: quant.uniform<i4:f32, 0.13329366842905679:-1>
  %0 = "quantfork.stats"(%arg0) {bitsNum = 8 : i64, layerStats = dense<[-0.999415695, 0.99998933]> : tensor<2xf32>, narrowRange = false} : (tensor<?x3xf32>) -> tensor<?x3xf32>
  // CHECK: %[[dot:.*]] = stablehlo.dot %[[dq2]], %[[dq1]]

// CHECK-NEXT:    "tf.MatMul"(%arg0, [[TENSOR]]) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
// CHECK-NEXT:    "tf.IfrtCall"(%arg0, [[ARRAYKEY]]) <{program_id = 6515870160938153680 : i64, variable_arg_indices = [1 : i32]}> {__tpu_compile_metadata_text = "retvals { sharding { } }"} : (tensor<1x3xf32>, tensor<!tf_type.string>) -> tensor<1x1xf32>
// CHECK-NEXT:    return
//

    %elem_shape_2 = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32>
    %tls = "tf.TensorListStack"(%tl_set_item, %elem_shape_2) {num_elements = 300 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<1xi32>) -> tensor<300x?xf32>

    %elem_2 = "tf.Const"() {value = dense<10.0> : tensor<9xf32>} : () -> tensor<9xf32>

    %1:1 = tf_executor.island wraps "tf.SendTPUEmbeddingGradients"(%0) {N = 1 : i64, NN = 0 : i64, config = "test_config_send_embedding", operandSegmentSizes = array<i32: 1, 0>} : (tensor<512x256xf32>) -> () loc("SendTPUEmbedding")
    tf_executor.fetch
  }
  func.return
}

// CHECK:       name: "RecvTPUEmbedding"