Mapping from logical cores in a computation to the physical TPU topology. Prefer to use the `DeviceAssignment.build()` helper to construct a `DeviceAssignment`; it is easier if less flexible than constructing a `DeviceAssignment` directly.
| 57 | |
| 58 | @tf_export("tpu.experimental.DeviceAssignment") |
| 59 | class DeviceAssignment(object): |
| 60 | """Mapping from logical cores in a computation to the physical TPU topology. |
| 61 | |
| 62 | Prefer to use the `DeviceAssignment.build()` helper to construct a |
| 63 | `DeviceAssignment`; it is easier if less flexible than constructing a |
| 64 | `DeviceAssignment` directly. |
| 65 | """ |
| 66 | |
| 67 | def __init__(self, topology, core_assignment): |
| 68 | """Constructs a `DeviceAssignment` object. |
| 69 | |
| 70 | Args: |
| 71 | topology: A `Topology` object that describes the physical TPU topology. |
| 72 | core_assignment: A logical to physical core mapping, represented as a |
| 73 | rank 3 numpy array. See the description of the `core_assignment` |
| 74 | property for more details. |
| 75 | |
| 76 | Raises: |
| 77 | ValueError: If `topology` is not `Topology` object. |
| 78 | ValueError: If `core_assignment` is not a rank 3 numpy array. |
| 79 | """ |
| 80 | if not isinstance(topology, Topology): |
| 81 | raise ValueError("topology must be a Topology object, got {}".format( |
| 82 | type(topology))) |
| 83 | core_assignment = np.asarray(core_assignment, dtype=np.int32) |
| 84 | |
| 85 | self._topology = topology |
| 86 | |
| 87 | if core_assignment.ndim != 3: |
| 88 | raise ValueError("core_assignment must be a rank 3 numpy array, " |
| 89 | "got shape {}".format(core_assignment.shape)) |
| 90 | |
| 91 | self._num_replicas = core_assignment.shape[0] |
| 92 | self._num_cores_per_replica = core_assignment.shape[1] |
| 93 | |
| 94 | if core_assignment.shape[-1] != topology.mesh_rank: |
| 95 | raise ValueError( |
| 96 | "minor dimension of core_assignment must have size equal to topology " |
| 97 | "rank ({}), got shape {}".format(topology.mesh_rank, |
| 98 | core_assignment.shape)) |
| 99 | |
| 100 | self._core_assignment = core_assignment |
| 101 | self._task_and_cores_to_replicas = _compute_task_and_cores_to_replicas( |
| 102 | self._core_assignment, topology) |
| 103 | |
| 104 | @property |
| 105 | def topology(self): |
| 106 | """A `Topology` that describes the TPU topology.""" |
| 107 | return self._topology |
| 108 | |
| 109 | @property |
| 110 | def num_cores_per_replica(self): |
| 111 | """The number of cores per replica.""" |
| 112 | return self._num_cores_per_replica |
| 113 | |
| 114 | @property |
| 115 | def num_replicas(self): |
| 116 | """The number of replicas of the computation.""" |