libinephany 0.15.2__tar.gz → 0.15.4__tar.gz

libinephany-0.15.4/CODE_VERSION.cfg

@@ -0,0 +1 @@
+0.15.4

{libinephany-0.15.2/libinephany.egg-info → libinephany-0.15.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: libinephany
-Version: 0.15.2
+Version: 0.15.4
 Summary: Inephany library containing code commonly used by multiple subpackages.
 Author-email: Inephany <info@inephany.com>
 License: Apache 2.0

{libinephany-0.15.2 → libinephany-0.15.4}/libinephany/observations/observation_utils.py

@@ -14,7 +14,7 @@ import torch
 import torch.optim as optim
 
 from libinephany.pydantic_models.schemas.tensor_statistics import TensorStatistics
-from libinephany.utils import optim_utils, torch_distributed_utils
+from libinephany.utils import optim_utils
 
 # ======================================================================================================================
 #
@@ -167,18 +167,39 @@ def create_one_hot_depth_encoding(agent_controlled_modules: list[str], parameter
     return create_one_hot_observation(vector_length=3, one_hot_index=one_hot_index)
 
 
+def check_if_tensor_is_xpu(tensor: torch.Tensor | None) -> bool:
+    """
+    :param tensor: Tensor to check whether it is an XPU tensor.
+    :return: Whether the tensor is an XPU tensor.
+    """
+
+    try:
+        xpu_available = torch.xpu.is_available()
+        return xpu_available and tensor.is_xpu
+
+    except AttributeError:
+        return False
+
+
 def tensor_on_local_rank(tensor: torch.Tensor | None) -> bool:
     """
     :param tensor: Tensor to check whether it is owned by the local rank, partially or entirely.
     :return: Whether the tensor is owned by the local rank.
     """
 
-    valid_tensor = tensor is not None and tensor.grad is not None and tensor.numel() > 0
+    valid_tensor = tensor is not None and tensor.numel() > 0
+    device_index = tensor.device.index if valid_tensor else None
+
+    is_xpu = check_if_tensor_is_xpu(tensor)
+    cuda_available = torch.cuda.is_available()
 
-    if valid_tensor and tensor.is_cuda:
-        local_rank = torch_distributed_utils.get_local_rank()
+    if valid_tensor and is_xpu:
+        current_device = torch.xpu.current_device()
+        return device_index == current_device
 
-        return tensor.device.index == local_rank
+    elif valid_tensor and cuda_available and tensor.is_cuda:
+        current_device = torch.cuda.current_device()
+        return device_index == current_device
 
     return valid_tensor
 

{libinephany-0.15.2 → libinephany-0.15.4}/libinephany/observations/statistic_trackers.py

@@ -403,7 +403,9 @@ class FirstOrderGradients(Statistic):
         :return: None, TensorStatistics model or a float.
         """
 
-        gradients = [p.grad.view(-1) for p in parameters if observation_utils.tensor_on_local_rank(p)]
+        gradients = [
+            p.grad.view(-1) for p in parameters if observation_utils.tensor_on_local_rank(p) and p.grad is not None
+        ]
 
         if not gradients:
             return None
@@ -495,7 +497,9 @@ class SecondOrderGradients(Statistic):
         :return: None, TensorStatistics model or a float.
         """
 
-        fo_gradients = [p.grad.view(-1) for p in parameters if observation_utils.tensor_on_local_rank(p)]
+        fo_gradients = [
+            p.grad.view(-1) for p in parameters if observation_utils.tensor_on_local_rank(p) and p.grad is not None
+        ]
 
         if not fo_gradients:
             return None

{libinephany-0.15.2 → libinephany-0.15.4/libinephany.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: libinephany
-Version: 0.15.2
+Version: 0.15.4
 Summary: Inephany library containing code commonly used by multiple subpackages.
 Author-email: Inephany <info@inephany.com>
 License: Apache 2.0

libinephany-0.15.2/CODE_VERSION.cfg

@@ -1 +0,0 @@
-0.15.2