(
self,
logical_mesh: Union[MeshShape, HybridMeshShape],
expected: Optional[Union[MeshShape, Exception]] = None,
)
| 1917 | def test_create_device_mesh_tpuv4( |
| 1918 | self, |
| 1919 | logical_mesh: Union[MeshShape, HybridMeshShape], |
| 1920 | expected: Optional[Union[MeshShape, Exception]] = None, |
| 1921 | ): |
| 1922 | physical_mesh = (4, 4, 1) |
| 1923 | coords = [ |
| 1924 | (x, y, z) |
| 1925 | for x in range(physical_mesh[0]) |
| 1926 | for y in range(physical_mesh[1]) |
| 1927 | for z in range(physical_mesh[2]) |
| 1928 | ] |
| 1929 | devices = [ |
| 1930 | DummyTpuDevice( |
| 1931 | platform="tpu", |
| 1932 | device_kind="TPU v4", |
| 1933 | process_index=ix // 4, |
| 1934 | coords=coord, |
| 1935 | ) |
| 1936 | for ix, coord in enumerate(coords) |
| 1937 | ] |
| 1938 | if isinstance(expected, Exception): |
| 1939 | with self.assertRaisesRegex(type(expected), str(expected)): |
| 1940 | create_device_mesh(mesh_shape=logical_mesh, devices=devices) |
| 1941 | else: |
| 1942 | # Check that the constructed mesh has the expected shape. |
| 1943 | self.assertEqual( |
| 1944 | expected or logical_mesh, |
nothing calls this directly
no test coverage detected