abstracting tpu type logic

Author: prishajain1

src/maxdiffusion/models/ltx2/attention_ltx2.py

@@ -20,6 +20,7 @@
 import jax.numpy as jnp
 from ... import common_types
 from ..attention_flax import NNXAttentionOp
+from maxdiffusion.tpu_utils import get_tpu_type, TpuType
 
 Array = common_types.Array
 Mesh = common_types.Mesh
@@ -360,8 +361,8 @@ def __init__(
     self.dropout_rate = dropout
 
     # Auto-detect hardware for sharding specs if not overridden
-    device_kind = jax.devices()[0].device_kind
-    is_ironwood = "7x" in device_kind
+    tpu_type = get_tpu_type()
+    is_ironwood = tpu_type == TpuType.TPU_7X
 
     if qkv_sharding_spec is None:
       qkv_sharding_spec = (None, "heads") if is_ironwood else ("embed", "heads")

src/maxdiffusion/tpu_utils.py

@@ -15,6 +15,7 @@
 """
 
 import jax
+from enum import Enum
 
 
 def print_device_memory_info(devices):
@@ -42,3 +43,23 @@ def print_array_info(array, name):
   for device_idx in num_devices:
     jax.debug.print("shape on device {x} : {y}", x=device_idx, y=array.device_buffers[0].shape)
     jax.debug.print("size on device {x} : {y}", x=device_idx, y=array.device_buffers[device_idx].size / array.size)
+
+
+class TpuType(Enum):
+  TPU_V6_LITE = "v6e"
+  TPU_7X = "v7x"
+  UNKNOWN = "unknown"
+
+
+def get_tpu_type() -> TpuType:
+  """Detects the current TPU hardware generation."""
+  try:
+    device_kind = jax.devices()[0].device_kind
+    if "7x" in device_kind:
+      return TpuType.TPU_7X
+    elif "v6 lite" in device_kind:
+      return TpuType.TPU_V6_LITE
+    else:
+      return TpuType.UNKNOWN
+  except Exception:
+    return TpuType.UNKNOWN