Searched refs:shardings (Results 1 – 14 of 14) sorted by relevance
72 OpSharding Tuple(const ShapeTree<OpSharding>& shardings) { in Tuple() argument75 for (const auto& index_to_sharding : shardings.leaves()) { in Tuple()
57 OpSharding Tuple(const ShapeTree<OpSharding>& shardings);
2 # Python API for shardings in XLA.
274 shardings = [278 type=xla_data_pb2.OpSharding.TUPLE, tuple_shardings=shardings)
444 const std::vector<HloSharding>& shardings) { in set_spmd_parameters_shardings() argument445 spmd_parameters_shardings_ = shardings; in set_spmd_parameters_shardings()
258 absl::Span<const HloSharding> shardings) { in Tuple() argument260 for (auto& sharding : shardings) { in Tuple()264 std::vector<HloSharding> flattened_list(shardings.begin(), shardings.end()); in Tuple()
97 absl::Span<const HloSharding> shardings);
40 // Whether to automatically generate XLA shardings for SPMD partitioner.
123 // Whether to automatically generate XLA shardings for SPMD partitioner.
703 // If type == TUPLE, the sub-shardings, one per leaf node in the tuple shape,706 // is not stored here; shardings do not store the shapes to which they are717 // combined with other shardings. Metadata are to not be populated when
493 // Whether to automatically generate XLA shardings for SPMD partitioner.
259 // enabled. Non replicated inputs/outputs should have shardings set to be282 // Tests partitioned variables (via XLA SPMD) propagates shardings correctly.
528 llvm::ArrayRef<std::optional<xla::OpSharding>> shardings) { in AllOptionalShardingsAreSet() argument529 return llvm::all_of(shardings, in AllOptionalShardingsAreSet()
6 // and op shardings are added. Sink tokens are created