Searched refs:SetSharding (Results 1 – 5 of 5) sorted by relevance
159 void SetSharding(const OpSharding& sharding) { sharding_ = sharding; } in SetSharding() function1057 SetSharding(sharding); in XlaScopedShardingAssignment()1064 ~XlaScopedShardingAssignment() { SetSharding(prev_sharding_); } in ~XlaScopedShardingAssignment()1067 void SetSharding(const absl::optional<OpSharding>& sharding) { in SetSharding() function1069 builder_->SetSharding(sharding.value()); in SetSharding()
748 def SetSharding(self, sharding): member in ComputationBuilder750 self._builder.SetSharding(sharding)
806 .def("SetSharding", &XlaBuilder::SetSharding) in PYBIND11_MODULE()
2046 c.SetSharding(sharding)
794 builder.SetSharding(sharding_builder::AssignDevice(1)); in XLA_TEST_F()