message: Also parse the public sharding attribute for XlaSharding op
  If this op is created from python API, the sharding attribute is missing.
diff: @@ -23,6 +23,7 @@ namespace tensorflow {
  namespace {
  const char kDeviceSuffixReplicatedCore[] = "REPLICATED_CORE";
  const char kShardingAttribute[] = "_XlaSharding";
  +const char kShardingOpAttribute[] = "sharding";
  } // namespace
  namespace {
  @@ -147,22 +148,42 @@ void SetShardingDeviceAssignmentFromNode(const Node& src, Node* dst) {
  }
  }
  -xla::StatusOr<absl::optional<xla::OpSharding>> GetShardingFromNodeDef(
  - const NodeDef& node_def, bool add_metadata) {
  - if (!HasNodeAttr(node_def, kShardingAttribute)) {
  +namespace {
  +
  +xla::StatusOr<absl::optional<xla::OpSharding>> GetShardingFromNodeDefInternal(
  + const NodeDef& node_def, bool add_metadata, const char* attribute) {
  + if (!HasNodeAttr(node_def, attribute)) {
  return absl::optional<xla::OpSharding>();
  }
  string value;
  xla::OpSharding sharding;
  - TF_RETURN_IF_ERROR(GetNodeAttr(node_def, kShardingAttribute, &value));
  + TF_RETURN_IF_ERROR(GetNodeAttr(node_def, attribute, &value));
  if (!sharding.ParseFromString(value)) {
  return xla::InvalidArgument(
  - "Experimental _XlaSharding attribute was not a valid encoded "
  - "xla::OpSharding proto.");
  + "Experimental %s attribute was not a valid encoded xla::OpSharding "
  + "proto.",
  + attribute);
  }
  if (add_metadata) {
  AssignOpMetadataToSharding(sharding, node_def.op(), node_def.name());
  }
  return absl::optional<xla::OpSharding>(sharding);
  }
  +
  +} // namespace
  +
  +xla::StatusOr<absl::optional<xla::OpSharding>> GetShardingFromNodeDef(
  + const NodeDef& node_def, bool add_metadata) {
  + if (node_def.op() == "XlaSharding") {
  + TF_ASSIGN_OR_RETURN(auto sharding,
  + GetShardingFromNodeDefInternal(node_def, add_metadata,
  + kShardingOpAttribute));
  + if (sharding.has_value()) {
  + return sharding;
  + }
  + }
  + return GetShardingFromNodeDefInternal(node_def, add_metadata,
  + kShardingAttribute);
  +}
  +
  } // namespace tensorflow
  