Searched refs:UsesDevice (Results 1 – 4 of 4) sorted by relevance
45 EXPECT_TRUE(sharding.UsesDevice(0)); in TEST_F()46 EXPECT_TRUE(sharding.UsesDevice(65535)); in TEST_F()60 EXPECT_FALSE(sharding.UsesDevice(0)); in TEST_F()61 EXPECT_TRUE(sharding.UsesDevice(5)); in TEST_F()
189 if (rhs.UsesDevice(kUnassignedDevice)) { in AssignLeafSharding()192 if (lhs->UsesDevice(kUnassignedDevice)) { in AssignLeafSharding()363 if (instruction->sharding().UsesDevice(kUnassignedDevice)) { in ApplyDomainSharding()
119 bool UsesDevice(int64 device) const;
108 bool HloSharding::UsesDevice(int64 device) const { in UsesDevice() function in xla::HloSharding111 return s.UsesDevice(device); in UsesDevice()