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

  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>>>,...

        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