Home
last modified time | relevance | path

Searched refs:device_assignment_lib (Results 1 – 5 of 5) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/python/distribute/
H A Dtpu_strategy_test.py58 from tensorflow.python.tpu import device_assignment as device_assignment_lib unknown
338 device_assignment = device_assignment_lib.DeviceAssignment.build(
345 device_assignment2 = device_assignment_lib.DeviceAssignment.build(
456 device_assignment = device_assignment_lib.DeviceAssignment.build(
464 device_assignment2 = device_assignment_lib.DeviceAssignment(
491 device_assignment = device_assignment_lib.DeviceAssignment.build(
499 device_assignment2 = device_assignment_lib.DeviceAssignment(
1401 device_assignment = device_assignment_lib.DeviceAssignment(
1413 device_assignment = device_assignment_lib.DeviceAssignment(
1426 device_assignment = device_assignment_lib.DeviceAssignment(
[all …]
H A Dstrategy_combinations.py40 from tensorflow.python.tpu import device_assignment as device_assignment_lib unknown
111 device_assignment = device_assignment_lib.DeviceAssignment(
113 core_assignment=device_assignment_lib.SINGLE_CORE_ASSIGNMENT)
H A Dtpu_strategy_model_parallelism_test.py42 from tensorflow.python.tpu import device_assignment as device_assignment_lib unknown
66 device_assignment = device_assignment_lib.DeviceAssignment.build(
H A Dtpu_strategy.py61 from tensorflow.python.tpu import device_assignment as device_assignment_lib # pylint: disable=unu… unknown
/aosp_15_r20/external/tensorflow/tensorflow/python/tpu/
H A Dtpu.py51 from tensorflow.python.tpu import device_assignment as device_assignment_lib unknown
237 ) -> Optional[device_assignment_lib.DeviceAssignment]:
919 device_assignment: Optional[device_assignment_lib.DeviceAssignment] = None,
1211 device_assignment: Optional[device_assignment_lib.DeviceAssignment] = None,
1788 device_assignment: Optional[device_assignment_lib.DeviceAssignment] = None,
1952 device_assignment: Optional[device_assignment_lib.DeviceAssignment] = None,
2036 device_assignment: Optional[device_assignment_lib.DeviceAssignment] = None,
2098 device_assignment: Optional[device_assignment_lib.DeviceAssignment] = None,
2255 device_assignment: Optional[device_assignment_lib.DeviceAssignment] = None,