Searched refs:UsesDevice (Results 1 – 6 of 6) sorted by relevance
189 if (rhs.UsesDevice(kUnassignedDevice)) { in AssignLeafSharding()192 if (lhs->UsesDevice(kUnassignedDevice)) { in AssignLeafSharding()363 if (instruction->sharding().UsesDevice(kUnassignedDevice)) { in ApplyDomainSharding()
60 EXPECT_TRUE(sharding.UsesDevice(0)); in TEST_F()61 EXPECT_TRUE(sharding.UsesDevice(65535)); in TEST_F()75 EXPECT_FALSE(sharding.UsesDevice(0)); in TEST_F()76 EXPECT_TRUE(sharding.UsesDevice(5)); in TEST_F()
159 bool UsesDevice(int64 device) const;
197 bool HloSharding::UsesDevice(int64 device) const { in UsesDevice() function in xla::HloSharding200 return s.UsesDevice(device); in UsesDevice()
173 if (element_sharding.UsesDevice(device)) { in GetPerDeviceShape()
195 if (sharding.UsesDevice(partition_id)) { in MakeNonPaddedShapeForGivenPartition()