Upload folder using huggingface_hub

modeling.py

@@ -5,12 +5,12 @@ import torch.nn.functional as F
 import math
 import importlib
 
-#try:
-#    xm = importlib.import_module('torch_xla.core.xla_model')
-#    xs = importlib.import_module('torch_xla.distributed.spmd.xla_sharding')
-#except ImportError:
-xm = None
-xs = None
+try:
+    xm = importlib.import_module('torch_xla.core.xla_model')
+    xs = importlib.import_module('torch_xla.distributed.spmd.xla_sharding')
+except ImportError:
+    xm = None
+    xs = None
 
 
 class Rotary3D(nn.Module):
@@ -102,8 +102,7 @@ class PSIAttentionLayer(nn.Module):
         # check if we are running on TPU
         try:
             # Use local import to avoid conflict if global xm is None and to check TPU specifically for this flag
-            #xm_local = importlib.import_module('torch_xla.core.xla_model')
-            raise ImportError
+            xm_local = importlib.import_module('torch_xla.core.xla_model')
             self.tpu = True
         except ImportError:
             self.tpu = False
@@ -153,7 +152,7 @@ class PSIAttentionLayer(nn.Module):
         # Apply attention
         if self.tpu:
             # (1)
-            #from torch_xla.experimental.custom_kernel import flash_attention
+            flash_attention = importlib.import_module('torch_xla.experimental.custom_kernel.flash_attention')
             q_norm = q / math.sqrt(k.size(-1))
             y = flash_attention(
                 q_norm, k, v, 

psi.py

@@ -18,12 +18,12 @@ from .modeling import (
     RMSNorm, PSIBlock, PartitionedEmbedding, PartitionedLinear
 )
 
-#try:
-#    xm = importlib.import_module('torch_xla.core.xla_model')
-#    xs = importlib.import_module('torch_xla.distributed.spmd.xla_sharding')
-#except ImportError:
-xm = None
-xs = None
+try:
+    xm = importlib.import_module('torch_xla.core.xla_model')
+    xs = importlib.import_module('torch_xla.distributed.spmd.xla_sharding')
+except ImportError:
+    xm = None
+    xs = None