if (node.IsWhileNode()) {
    auto* output_shapes = node.attrs().Find("output_shapes");
    auto* element_types = node.attrs().Find("T");
    if (output_shapes && !output_shapes->list().shape().empty()) {
      const auto& output_shape = output_shapes->list().shape(idx);
      const auto& element_type = element_types->list().type(idx);
      return ConvertToMlirTensorType(output_shape, element_type, &builder);
    }
  }

    %0 = "tf.TakeDataset"(%arg0, %cst) {device = "", metadata = "", output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]} : (tensor<!tf_type.variant>, tensor<i64>) -> tensor<!tf_type.variant>

        // expected-remark@above {{ID: 0}}

  %4 = "tf.Case"(%2, %arg0, %arg1) {branches = [@sub, @add], output_shapes = [#tf_type.shape<>], device = "noodle", is_stateless = false} : (tensor<i32>, tensor<f32>, tensor<f32>) -> tensor<f32>
  // CHECK: PartitionedCall
  // CHECK-SAME: f = @sub
  // CHECK-SAME: _cluster_launch = "not_ready"
  // CHECK-SAME: device = "noodle"

  }

  auto output_shape = xla::ShapeInference::InferGatherShape(
      input_shape, start_indices_shape, gather_dim_numbers, slice_sizes);
  if (!output_shape.ok()) {
    op->emitError() << output_shape.status().message();
    return false;
  }

  auto refined_type = xla::ConvertShapeToType<RankedTensorType>(
      *output_shape, mlir::Builder(op));
  if (!refined_type.ok()) {

      %arg1: tensor<i64>
    ) {
      %1 = "tf.ReduceDataset"(%arg0, %arg1) {
        Targuments = [],
        Tstate = [i64], device = "",
        f = @__reduce_func_1, f._tf_data_function = true,
        output_shapes = [#tf_type.shape<>],
        output_types = [i64], use_inter_op_parallelism = true, _xla_compile_device_type="TPU"} :
          (tensor<!tf_type.variant>, tensor<i64>) -> (tensor<i64>)
      func.return
    }
    ```

  int64_t out_row_dim = output_shape[output_shape.size() - 2];
  int64_t out_col_dim = output_shape[output_shape.size() - 1];

  int64_t expected_out_row_dim = op.getAdjX() ? x_col_dim : x_row_dim;
  int64_t expected_out_col_dim = op.getAdjY() ? y_row_dim : y_col_dim;

  if (expected_out_row_dim != ShapedType::kDynamic &&
      out_row_dim != ShapedType::kDynamic &&
      out_row_dim != expected_out_row_dim)

  %weight_scales: tensor<f32>, %weight_zps: tensor<i32>,
  %output_scales: tensor<f32>, %output_zps: tensor<i32>) -> () {
  // expected-error @below {{'tf.UniformQuantizedDot' op quantization_axis is -1, scales must have 0 rank.}}
  %1 = "tf.UniformQuantizedDot"(
    %input, %weight,
    %input_scales, %input_zps,
    %weight_scales, %weight_zps,
    %output_scales, %output_zps) {
      lhs_quantization_axis = -1 : i64,

  %4 = "tf.Add"(%2, %arg2) : (tensor<?x8x10xf32>, tensor<8x10xf32>) -> tensor<?x8x10xf32>
  %5 = "tf.Add"(%arg1, %arg2) : (tensor<8x10xf32>, tensor<8x10xf32>) -> tensor<8x10xf32>
  %6 = "tf.Const"() {_output_shapes = ["tfshape$"], device = "/device:CPU:0", dtype = f32, value = dense<1.000000e+00> : tensor<f32>} : () -> tensor<f32>

  int64_t out_row_dim = output_shape[output_shape.size() - 2];
  int64_t out_col_dim = output_shape[output_shape.size() - 1];

  int64_t expected_out_row_dim = op.getAdjX() ? x_col_dim : x_row_dim;
  int64_t expected_out_col_dim = op.getAdjY() ? y_row_dim : y_col_dim;

  if (expected_out_row_dim != ShapedType::kDynamic &&
      out_row_dim != ShapedType::kDynamic &&
      out_row_dim != expected_out_row_dim)