Searched refs:tuple_shardings (Results 1 – 4 of 4) sorted by relevance
129 tuple_shardings = list(proto.tuple_shardings)130 tuple_shardings[tensor.value_index] = self._proto132 type=xla_data_pb2.OpSharding.TUPLE, tuple_shardings=tuple_shardings)161 type=xla_data_pb2.OpSharding.TUPLE, tuple_shardings=shardings)
363 std::vector<HloSharding> tuple_shardings; in FromProto() local364 tuple_shardings.reserve(proto.tuple_shardings().size()); in FromProto()365 for (const OpSharding& tuple_sharding_proto : proto.tuple_shardings()) { in FromProto()368 tuple_shardings.push_back(sharding); in FromProto()370 return HloSharding(tuple_shardings); in FromProto()
240 explicit HloSharding(const std::vector<HloSharding>& tuple_shardings) in HloSharding() argument245 tuple_elements_(tuple_shardings) {}
582 // This sharding is a tuple - only the tuple_shardings field is valid.603 repeated OpSharding tuple_shardings = 5; field