Merge pull request tensorflow#59998 from tensorflow/fix-bad-cherrypick-again

Fix bad cherrypick

Author: mihaimaruseac

tensorflow/core/tpu/ops/tpu_partitioned_input_op.cc

@@ -124,85 +124,4 @@ REGISTER_OP("TPUPartitionedInput")
       return OkStatus();
     });
 
-REGISTER_OP("TPUPartitionedInputV2")
-    .Input("inputs: N * T")
-    .Output("output: T")
-    .Attr("N: int >= 1")
-    .Attr("T: type")
-    .Attr("partition_dims: list(int)")
-    .Attr("is_packed: bool = false")
-    .SetShapeFn([](InferenceContext* c) {
-      DataType dtype;
-      TF_RETURN_IF_ERROR(c->GetAttr("T", &dtype));
-      std::vector<int> partition_dims;
-      TF_RETURN_IF_ERROR(c->GetAttr("partition_dims", &partition_dims));
-      bool is_packed;
-      TF_RETURN_IF_ERROR(c->GetAttr("is_packed", &is_packed));
-
-      int num_partitions = 1;
-      for (const int& partition_dim : partition_dims) {
-        num_partitions *= partition_dim;
-      }
-
-      bool replicated = partition_dims.empty();
-      int num_inputs_expected = is_packed ? 1 : num_partitions;
-      if (!((replicated && !is_packed) ||
-            (c->num_inputs() == num_inputs_expected))) {
-        // we cannot validate the number of inputs for replicated, unpacked ops
-        // since we cannot infer the number of partitions from partition_dims
-        return errors::InvalidArgument("Expected ", num_inputs_expected,
-                                       " inputs, got ", c->num_inputs(), ".");
-      } else if (c->num_inputs() == 0) {
-        return errors::InvalidArgument(
-            "Expected at least one input to TPUPartitionedInputV2.");
-      }
-
-      ShapeHandle output_shape;
-      if (dtype == DT_RESOURCE) {
-        ShapeHandle previous_shape_handle;
-        const std::vector<shape_inference::ShapeAndType>* shapes_and_types =
-            nullptr;
-        for (int i = c->num_inputs() - 1; i >= 0; --i) {
-          shapes_and_types = c->input_handle_shapes_and_types(i);
-          if (shapes_and_types) {
-            ShapeHandle shape_handle = shapes_and_types->at(0).shape;
-            if (!c->FullyDefined(shape_handle)) {
-              return errors::InvalidArgument("Inputs must have static shape,",
-                                             "input[", i,
-                                             "] has unknown dimension.");
-            }
-
-            if (i != c->num_inputs() - 1) {
-              ShapeHandle tmp;
-              if (!c->Merge(shape_handle, previous_shape_handle, &tmp).ok()) {
-                return errors::InvalidArgument(
-                    "Inputs must have the same shape.");
-              }
-            } else {
-              previous_shape_handle = shape_handle;
-            }
-          }
-        }
-
-        if (shapes_and_types) {
-          TF_ASSIGN_OR_RETURN(
-              output_shape,
-              _ComputeOutputShape(c, previous_shape_handle, partition_dims));
-          std::vector<shape_inference::ShapeAndType> output_shapes_and_types;
-          output_shapes_and_types.push_back(shape_inference::ShapeAndType(
-              output_shape, shapes_and_types->at(0).dtype));
-          c->set_output_handle_shapes_and_types(0, output_shapes_and_types);
-        }
-      }
-
-      if (!c->FullyDefined(output_shape)) {
-        TF_ASSIGN_OR_RETURN(
-            output_shape, _ComputeOutputShape(c, c->input(0), partition_dims));
-      }
-
-      c->set_output(0, output_shape);
-
-      return OkStatus();
-    });
-
 }  // namespace tensorflow

tensorflow/core/tpu/ops/tpu_partitioned_output_op.cc

@@ -70,55 +70,4 @@ REGISTER_OP("TPUPartitionedOutput")
       return OkStatus();
     });
 
-REGISTER_OP("TPUPartitionedOutputV2")
-    .Input("inputs:  T")
-    .Output("output: num_splits * T")
-    .Attr("T: type")
-    .Attr("num_splits: int >= 1")
-    .Attr("partition_dims: list(int)")
-    .SetShapeFn([](InferenceContext* c) {
-      DataType dtype;
-      TF_RETURN_IF_ERROR(c->GetAttr("T", &dtype));
-      std::vector<int> partition_dims;
-      TF_RETURN_IF_ERROR(c->GetAttr("partition_dims", &partition_dims));
-      int num_splits;
-      TF_RETURN_IF_ERROR(c->GetAttr("num_splits", &num_splits));
-      if (dtype == DT_RESOURCE) {
-        return errors::Unimplemented("Not implemented.");
-      } else if (c->num_inputs() == 0) {
-        return errors::InvalidArgument(
-            "Expected at least one input to TPUPartitionedOutputV2.");
-      }
-
-      ShapeHandle handle = c->input(0);
-      int rank = InferenceContext::Rank(handle);
-      int num_cores_per_replica = 1;
-      for (const int& partition_dim : partition_dims) {
-        num_cores_per_replica *= partition_dim;
-      }
-
-      if (num_splits != num_cores_per_replica) {
-        return errors::InvalidArgument("Expected ", num_cores_per_replica,
-                                       " splits.");
-      } else if (rank > (int)partition_dims.size()) {
-        return errors::InvalidArgument("Expected at least ", rank,
-                                       " partition dimensions.");
-      }
-
-      for (int i = 0; i < rank; ++i) {
-        shape_inference::DimensionHandle dim;
-        TF_RETURN_WITH_CONTEXT_IF_ERROR(
-            c->Divide(c->Dim(handle, i), partition_dims[i],
-                      true /* evenly_divisible */, &dim),
-            "Number of ways to split should evenly divide the split dimension");
-        TF_CHECK_OK(c->ReplaceDim(handle, i, dim, &handle));
-      }
-
-      for (int i = num_splits - 1; i >= 0; --i) {
-        c->set_output(i, handle);
-      }
-
-      return OkStatus();
-    });
-
 }  // namespace tensorflow