// dimensions.
  SmallVector<int64_t> output_shape(output_tensor_rank);
  for (int i = 0; i < output_tensor_rank; i++) {
    if (collapsed_dims.contains(i)) {
      // The collapsed dimension's size should have been 1, so it restores the
      // dimension with size 1.
      output_shape[i] = 1;
    } else {
      output_shape[i] = *shape_itr;
      shape_itr++;
    }
  }

  %bias_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
  %bias_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>

  %output_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
  %output_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>

    tf_executor.fetch
  }
  func.return
}

// CHECK:       name: "RecvTPUEmbedding"
// CHECK-NEXT:  op: "RecvTPUEmbeddingActivations"
// CHECK-NEXT:  attr {
// CHECK-NEXT:    key: "_output_shapes"
// CHECK-NEXT:    value {
// CHECK-NEXT:      list {
// CHECK-NEXT:        shape {
// CHECK-NEXT:          dim {
// CHECK-NEXT:            size: 512
// CHECK-NEXT:          }
// CHECK-NEXT:          dim {

  SmallVector<int64_t> output_shape(input_shape.getShape().begin(),
                                    input_shape.getShape().end());
  for (int i : spatial_dims) {
    output_shape[i] += padding_values[2 * i] + padding_values[2 * i + 1];
  }

  return builder.create<TF::PadV2Op>(
      loc, RankedTensorType::get(output_shape, builder.getI8Type()), input,
      temp_padding,

    %arg1: tensor<1x3xf32> {tf_saved_model.index_path = [0]})
      -> (tensor<3x3xf32> {tf_saved_model.index_path = []}) {
  // CHECK: {{%.*}} = corert.get_op_handler %arg0 "/device:GPU:0"
  %2 = "tf.MatMul"(%arg0, %arg1) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "/device:GPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
  func.return %2 : tensor<3x3xf32>

    const mlir::TensorType cluster_func_output_type,
    const xla::OpSharding& output_sharding,
    mlir::Type* tiled_logical_computation_type) {
  const auto output_shape = cluster_func_output_type.getShape();
  auto new_output_shape = llvm::to_vector<4>(output_shape);
  auto dimension_to_splits_map =
      GetDimensionIndicesAndNumSplitsFromSharding(output_sharding);
  if (!dimension_to_splits_map.ok()) {

  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,

// Converts a shape from MLIR to a TensorFlow tensor shape proto.
void ConvertToTensorShapeProto(llvm::ArrayRef<int64_t> shape,
                               TensorShapeProto* output_shape);

// Converts an MLIR type to a TensorFlow tensor shape.
PartialTensorShape ConvertTypeToTensorShape(const mlir::Type& type);

// Converts an MLIR shaped type to a TensorFlow shape attribute.

  }

  // Gather shapes for output.
  for (auto v : ddn.lhs_batch_dimensions()) {
    output_shape.push_back(lhs_shape[v]);
  }

  // Batch dimension is gathered from the right side.
  if (output_shape.empty()) {
    for (auto v : ddn.rhs_batch_dimensions()) {
      output_shape.push_back(rhs_shape[v]);
    }
  }

  // Gather remaining dimensions.