def from_str(layout_str: bytes) -> 'Layout':
layout_proto = layout_pb2.LayoutProto()
layout_proto.ParseFromString(layout_str)
sharding_specs = [
sharding_spec.sharding_spec
for sharding_spec in layout_proto.sharding_specs
]
mesh = Mesh.from_proto(layout_proto.mesh_config)
return Layout(sharding_specs, mesh)
def from_str(layout_str: bytes) -> 'Layout':
"""Creates an instance from a serialized Protobuf binary string."""
layout_proto = layout_pb2.LayoutProto()
layout_proto.ParseFromString(layout_str)
sharding_specs = [
sharding_spec.sharding_spec
for sharding_spec in layout_proto.sharding_specs
]
mesh = Mesh.from_proto(layout_proto.mesh_config)
return Layout(sharding_specs, mesh)
def as_proto(self) -> layout_pb2.LayoutProto:
"""Create a proto representation of a layout."""
layout_proto = layout_pb2.LayoutProto()
for dim_sharding in self.sharding_specs:
tensor_dim = layout_proto.sharding_specs.add()
tensor_dim.sharding_spec = dim_sharding
layout_proto.mesh_config.CopyFrom(self.mesh_proto())
return layout_proto