"tf.Yield"(%1) : (tensor<5xf32>) -> ()
  }) {device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0", metadata = "", operandSegmentSizes = array<i32: 1, 2, 1>, output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]} : (tensor<4xf32>, tensor<3xf32>, tensor<!tf_type.resource>, tensor<2xf32>) -> tensor<!tf_type.variant>
  return
}

// -----

    // Shape of the final output. (Except for dimension folding in the
    // single diagonal case.)
    Shape output_shape;
    for (int i = 0; i < num_dims - 2; i++) {
      output_shape.push_back(input_type.getDimSize(i));
    }
    output_shape.push_back(num_diags);
    output_shape.push_back(max_diag_len);

    // A slice is the shape of what GatherOp copies per lookup. So the last

      %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
    }
    ```

  // CHECK-DAG: [[OUTPUT_SHAPE:%.+]] = "tf.ConcatV2"([[OUTPUT_BATCH]], [[OUTER_SHAPE_0]], [[OUTER_SHAPE_1]], [[PADDED_SHAPE_SPLITS]]#3, [[ZERO_I64]])
  // CHECK-DAG: [[RESHAPED:%.+]] = "tf.Reshape"([[PADDED]], [[RESHAPED_SHAPE]])
  // CHECK-DAG: [[PERMUTED:%.+]] = "tf.Transpose"([[RESHAPED]], [[PERMUTATION]])
  // CHECK-DAG: [[RESULT:%.+]] = "tf.Reshape"([[PERMUTED]], [[OUTPUT_SHAPE]])
  // CHECK-DAG: return [[RESULT]]

    %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
 }
 ```

  let summary = "Transpose convolution operator";

  let description = [{
    Performs transpose convolution operation on input.
  }];

  let arguments = (ins
    TFL_I32Tensor:$output_shape,
    TFL_TensorOf<[F32, QI8, QUI8, QI16]>:$weights,
    TFL_TensorOf<[F32, QI8, QUI8, QI16]>:$input,
    TFL_TensorOfOrNone<[F32, QI32, I64]>:$bias,
    TFL_PaddingAttr:$padding,

                    IsWithinInt32Range);

    if (elements_all_in_int32_range) {
      std::vector<int32_t> output_shape(output_ty.getRank());
      std::transform(output_ty.getShape().begin(), output_ty.getShape().end(),
                     output_shape.begin(),
                     [](int64_t val) { return static_cast<int32_t>(val); });
      output_int_type = tensorflow::GetTypeFromTFTensorShape(

[i32, i32, i32, !tf_type.variant, !tf_type.variant, f32], _lower_using_switch_merge = true, _num_original_outputs = 6 : i64, _read_only_resource_inputs = [], body = @map_while_body_180, cond = @map_while_cond_170, device = "", is_stateless = true, output_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<?>], parallel_iterations = 4 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>,...

      device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0",
      finalize_func = @finalize,
      init_func = @init,
      next_func = @next,
      operandSegmentSizes = array<i32: 1, 2, 1>,
      output_shapes = [#tf_type.shape<>],
      output_types = [!tf_type.string],
      metadata = ""} : (
              tensor<4xf32>,
              tensor<3xf32>,
              tensor<!tf_type.resource>,

        all the TensorFlow OPs included in each of the `TRTEngineOp`s.

*   `tf.tpu.experimental.embedding`:

    *   `tf.tpu.experimental.embedding.FeatureConfig` now takes an additional
        argument `output_shape` which can specify the shape of the output
        activation for the feature.
    *   `tf.tpu.experimental.embedding.TPUEmbedding` now has the same behavior