Searched refs:experimental_assign_to_logical_device (Results 1 – 3 of 3) sorted by relevance
36 name: "experimental_assign_to_logical_device"
177 strategy.experimental_assign_to_logical_device(x, 0)
441 def experimental_assign_to_logical_device(self, tensor, logical_device_id): member in TPUStrategyV2