libinephany 0.15.1__py3-none-any.whl → 0.15.3__py3-none-any.whl

libinephany/observations/observation_utils.py

@@ -173,7 +173,21 @@ def tensor_on_local_rank(tensor: torch.Tensor | None) -> bool:
     :return: Whether the tensor is owned by the local rank.
     """
 
-    return 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
+
+    xpu_available = torch.xpu.is_available()
+    cuda_available = torch.cuda.is_available()
+
+    if valid_tensor and xpu_available and tensor.is_xpu:
+        current_device = torch.xpu.current_device()
+        return device_index == current_device
+
+    elif valid_tensor and cuda_available and tensor.is_cuda:
+        current_device = torch.cuda.current_device()
+        return device_index == current_device
+
+    return valid_tensor
 
 
 def form_update_tensor(

libinephany/observations/statistic_trackers.py

@@ -193,13 +193,14 @@ class Statistic(ABC):
         Processes the tensor cache to build a TensorStatistic model.
         """
 
-        concatenated = torch.cat(self._tensor_cache)
-        self._tensor_cache = []
+        if self._tensor_cache:
+            concatenated = torch.cat(self._tensor_cache)
+            self._tensor_cache = []
 
-        statistics = TensorStatistics.build(
-            tensor=concatenated, skip_statistics=self.skip_statistics, sample_percentage=self.downsample_percent
-        )
-        self._data.append(statistics)  # type: ignore
+            statistics = TensorStatistics.build(
+                tensor=concatenated, skip_statistics=self.skip_statistics, sample_percentage=self.downsample_percent
+            )
+            self._data.append(statistics)  # type: ignore
 
     @staticmethod
     @final
@@ -281,7 +282,6 @@ class Statistic(ABC):
 
         parameter_group = self._find_parameter_group(optimizer=optimizer)
         parameters = self._get_parameters(parameter_group=parameter_group)
-        self._sample_number += 1
 
         if self._sample_number % self.sample_frequency == 0:
             statistic = self._gather(
@@ -290,18 +290,18 @@ class Statistic(ABC):
 
             statistic = self._distributed_reduce(statistic=statistic)
 
-            if not torch_distributed_utils.is_scheduler_master_rank():
-                return
+            if torch_distributed_utils.is_scheduler_master_rank():
+                if isinstance(statistic, torch.Tensor):
+                    statistic = statistic.view(-1)
+                    self._tensor_cache.append(statistic)
 
-            if isinstance(statistic, torch.Tensor):
-                statistic = statistic.view(-1)
-                self._tensor_cache.append(statistic)
+                    if len(self._tensor_cache) >= self.max_cache_size:
+                        self._process_tensor_cache()
 
-                if len(self._tensor_cache) >= self.max_cache_size:
-                    self._process_tensor_cache()
+                elif statistic is not None:
+                    self._data.append(statistic)  # type: ignore
 
-            elif statistic is not None:
-                self._data.append(statistic)  # type: ignore
+        self._sample_number += 1
 
     @final
     def fetch(self) -> TensorStatistics | float | None:
@@ -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.1.dist-info → libinephany-0.15.3.dist-info}/METADATA

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

{libinephany-0.15.1.dist-info → libinephany-0.15.3.dist-info}/RECORD

@@ -2,11 +2,11 @@ libinephany/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 libinephany/aws/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 libinephany/aws/s3_functions.py,sha256=W8u85A6tDloo4FlJvydJbVHCUq_m9i8KDGdnKzy-Xpg,1745
 libinephany/observations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-libinephany/observations/observation_utils.py,sha256=wb6EZiaEiPuOqN26zzuT1rHyehoKh-8818KXn8pHweI,8688
+libinephany/observations/observation_utils.py,sha256=IidbKlZ4YaLnUFwMFcUdHVQjxEu78bvpD_nqaCnrgTA,9168
 libinephany/observations/observer_pipeline.py,sha256=ZhONGXJQSgs2VJJn9d2F7ItkYqntvchl9-JTyxW9eU0,12146
 libinephany/observations/pipeline_coordinator.py,sha256=FrN3linKaC0pVE5uKjlh_0Fi8Mb1oK91NzH3Fq7PvyM,7420
 libinephany/observations/statistic_manager.py,sha256=LLg1zSxnJr2oQQepYla3qoUuRy10rsthr9jta4wEbnc,8956
-libinephany/observations/statistic_trackers.py,sha256=J444i9EZ30vcYOEYqcDBzz7_UDpEE2hW_ISYBu_hwYc,30180
+libinephany/observations/statistic_trackers.py,sha256=PUBqGgMRi51SmiNh5HAH5kpYxsaflRepmM-uKyMiQZg,30326
 libinephany/observations/observers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 libinephany/observations/observers/base_observers.py,sha256=RkG5SW0b6Ooy0_oscRHxyB_YFNP7k8fxu37jBZElxIM,15418
 libinephany/observations/observers/global_observers.py,sha256=3TaiV2AxMOXfDq-kXMU3ZSo-rQENNCFhdWCJtpY99ok,38684
@@ -50,8 +50,8 @@ libinephany/utils/typing.py,sha256=rGbaPO3MaUndsWiC_wHzReD_TOLYqb43i01pKN-j7Xs,6
 libinephany/web_apps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 libinephany/web_apps/error_logger.py,sha256=gAQIaqerqP4ornXZwFF1cghjnd2mMZEt3aVrTuUCr34,16653
 libinephany/web_apps/web_app_utils.py,sha256=qiq_lasPipgN1RgRudPJc342kYci8O_4RqppxmIX8NY,4095
-libinephany-0.15.1.dist-info/licenses/LICENSE,sha256=pogfDoMBP07ehIOvWymuWIar8pg2YLUhqOHsJQU3wdc,9250
-libinephany-0.15.1.dist-info/METADATA,sha256=lsqYtqyJ_k_clascJkzx8rR7gEN75tZ8lCKNqcH1cps,8354
-libinephany-0.15.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-libinephany-0.15.1.dist-info/top_level.txt,sha256=bYAOXQdJgIoLkO2Ui0kxe7pSYegS_e38u0dMscd7COQ,12
-libinephany-0.15.1.dist-info/RECORD,,
+libinephany-0.15.3.dist-info/licenses/LICENSE,sha256=pogfDoMBP07ehIOvWymuWIar8pg2YLUhqOHsJQU3wdc,9250
+libinephany-0.15.3.dist-info/METADATA,sha256=UuSiJtqLBermVKve6SNjBQz_ViSvMVf5NEootGJNp4w,8354
+libinephany-0.15.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+libinephany-0.15.3.dist-info/top_level.txt,sha256=bYAOXQdJgIoLkO2Ui0kxe7pSYegS_e38u0dMscd7COQ,12
+libinephany-0.15.3.dist-info/RECORD,,