message: Add support for propagating resource shapes via the TPUReplicatedInput operator's shape inference function.
diff: @@ -44,6 +44,27 @@ REGISTER_OP("TPUReplicatedInput")
  " with other shapes.");
  }
  c->set_output(0, cur);
  +
  + // If this is a resource, unify the resource shapes.
  + DataType dtype;
  + TF_RETURN_IF_ERROR(c->GetAttr("T", &dtype));
  + if (dtype == DT_RESOURCE) {
  + const std::vector<shape_inference::ShapeAndType>* shapes_and_types =
  + nullptr;
  + for (int i = c->num_inputs() - 1; i >= 0; --i) {
  + if (shapes_and_types) {
  + if (!c->MergeInputHandleShapesAndTypes(i, *shapes_and_types)) {
  + return errors::InvalidArgument(
  + "Incompatible resource shapes for replicated TPU input.");
  + }
  + } else {
  + shapes_and_types = c->input_handle_shapes_and_types(i);
  + }
  + }
  + if (shapes_and_types) {
  + c->set_output_handle_shapes_and_types(0, *shapes_and_types);
  + }
  + }
  return Status::OK();
  })
  .Doc(
  