Searched refs:SetSharding (Results 1 – 4 of 4) sorted by relevance
405 .def("set_sharding", &XlaBuilder::SetSharding) in BuildXlaCompilerSubmodule()
194 void SetSharding(const OpSharding& sharding) { sharding_ = sharding; } in SetSharding() function1485 SetSharding(sharding); in XlaScopedShardingAssignment()1492 ~XlaScopedShardingAssignment() { SetSharding(prev_sharding_); } in ~XlaScopedShardingAssignment()1495 void SetSharding(const absl::optional<OpSharding>& sharding) { in SetSharding() function1497 builder_->SetSharding(sharding.value()); in SetSharding()
1336 builder->SetSharding(sharding); in Lower()1501 builder->SetSharding(sharding); in SetEntryTupleShardings()1539 builder->SetSharding(*arg_shardings[arg.getArgNumber()]); in LowerBasicBlockAsFunction()
794 builder.SetSharding(sharding_builder::AssignDevice(1)); in XLA_TEST_F()