This quick change folds the additional pf dimension in the topology into the 1st mesh tf dimension.

Folding the additional pf dimension into the 1st and 2nd mesh tf dimensions is equivalent w.r.t. examples and global_step per sec.

PiperOrigin-RevId: 355756187

Author: Mesh TensorFlow Team

mesh_tensorflow/simd_mesh_impl.py

@@ -826,9 +826,9 @@ def physical_shape_3d_from_topology_proto_4d(mesh_shape):
   Returns:
     a list of length 3
   """
-  if len(mesh_shape) != 4 or mesh_shape[2] != 1:
-    raise ValueError("Expected a 4d shape [x, y, 1, core]")
-  return [mesh_shape[1], mesh_shape[0], mesh_shape[3]]
+  if len(mesh_shape) != 4:
+    raise ValueError("Expected a 4d shape [x, y, z, core]")
+  return [mesh_shape[1]*mesh_shape[2], mesh_shape[0], mesh_shape[3]]
 
 
 def auto_logical_to_physical_tpu(logical_shape,

mesh_tensorflow/transformer/utils.py

@@ -255,8 +255,8 @@ def tpu_mesh_shape(tpu_topology=gin.REQUIRED,
   if tpu_topology.startswith("v"):
     num_cores = int(tpu_topology.split("-")[-1])
   else:
-    x, y = tpu_topology.split("x")
-    num_cores = int(x) * int(y) * 2
+    tpu_dim = [int(x) for x in tpu_topology.split("x")]
+    num_cores = functools.reduce(lambda x, y: x*y, tpu_dim) * 2
   if isinstance(model_parallelism, list):
     # model_parallelism is actually a spec used to
     # construct a simd_mesh_impl.HierarchicalTiling object