nshtrainer 1.3.0__py3-none-any.whl → 1.3.2__py3-none-any.whl

nshtrainer/model/base.py

@@ -85,12 +85,25 @@ def default_split_batched_predictions(
     """
     import torch.utils._pytree as tree
 
-    for sample_idx, batch_idx in enumerate(batch_indices):
+    for i, global_idx in enumerate(batch_indices):
+
+        def _verify_and_index(x: torch.Tensor):
+            # Make sure dim 0 length is equal to the batch size,
+            # otherwise we can't index it and should prompt
+            # the user to implement a splitter
+            if x.shape[0] != len(batch_indices):
+                raise ValueError(
+                    f"Batch size {x.shape[0]} does not match the number of batch indices {len(batch_indices)}. "
+                    "Please implement a custom `split_batched_predictions` method in your LightningModuleBase class."
+                )
+
+            return x[i]
+
         # Create a dictionary for each sample
         yield IndividualSample(
-            index=batch_idx,
-            batch=tree.tree_map(lambda x: x[sample_idx], batch),
-            prediction=tree.tree_map(lambda x: x[sample_idx], prediction),
+            index=global_idx,
+            batch=tree.tree_map(_verify_and_index, batch),
+            prediction=tree.tree_map(_verify_and_index, prediction),
         )
 
 

nshtrainer/trainer/__init__.py

@@ -1,7 +1,13 @@
 from __future__ import annotations
 
 from ..callbacks import callback_registry as callback_registry
+from ..callbacks.distributed_prediction_writer import (
+    DistributedPredictionReader as DistributedPredictionReader,
+)
 from ._config import TrainerConfig as TrainerConfig
+from ._distributed_prediction_result import (
+    DistributedPredictionResult as DistributedPredictionResult,
+)
 from .accelerator import accelerator_registry as accelerator_registry
 from .plugin import plugin_registry as plugin_registry
 from .trainer import Trainer as Trainer

{nshtrainer-1.3.0.dist-info → nshtrainer-1.3.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.3.0
+Version: 1.3.2
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.3.0.dist-info → nshtrainer-1.3.2.dist-info}/RECORD

@@ -118,7 +118,7 @@ nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=irPyDjfUX843ze4bJM9sW8WSe
 nshtrainer/metrics/__init__.py,sha256=Nqkn_jsDf3n5WtfMcnaaEftYjIIT2b-S7rmsB1MOMkU,86
 nshtrainer/metrics/_config.py,sha256=ox_ScK6V0J9nzIMhEB0qpToNKpt83VVgOVSRFCV-wBc,595
 nshtrainer/model/__init__.py,sha256=3G-bwPPSRStWdsdwG9-rn0bXcRpEiP1BiQpF_qavtls,97
-nshtrainer/model/base.py,sha256=Pv3M3QStWQp-DnfGFsLPAmp87HHrX1NrkAa4JcyBoDk,10255
+nshtrainer/model/base.py,sha256=PvTmupfGahEZME0BWqbeErDPP1VOm2Nm9JxJkO8afcc,10815
 nshtrainer/model/mixins/callback.py,sha256=0LPgve4VszHbLipid4mpI1qnnmdGS2spivs0dXLvqHw,3154
 nshtrainer/model/mixins/debug.py,sha256=ydLuAAaa7M5bX0gougZ5gWuZnvn4Ra9assal3IZ9hq8,2086
 nshtrainer/model/mixins/logger.py,sha256=7u9fQig-SVFA9RFIB4U0gqJAzruh49mgmXXvZ6VkDUk,11694
@@ -134,7 +134,7 @@ nshtrainer/profiler/_base.py,sha256=kFcSVn9gJuMwgDxbfyHh46CmEAIPZjxw3yjPbKgzvwA,
 nshtrainer/profiler/advanced.py,sha256=XrM3FX0ThCv5UwUrrH0l4Ow4LGAtpiBww2N8QAU5NOQ,1160
 nshtrainer/profiler/pytorch.py,sha256=8K37XvPnCApUpIK8tA2zNMFIaIiTLSoxKQoiyCPBm1Q,2757
 nshtrainer/profiler/simple.py,sha256=PimjqcU-JuS-8C0ZGHAdwCxgNLij4x0FH6WXsjBQzZs,1005
-nshtrainer/trainer/__init__.py,sha256=fQ7gQRlGWX-90TYT0rttkQyvXDCzo7DAvJgr-jX1zsY,316
+nshtrainer/trainer/__init__.py,sha256=jRaHdaFK8wxNrN1bleT9cf29iZahL_-XkWo5TWz2CmA,550
 nshtrainer/trainer/_config.py,sha256=Lt9tuzxgVzVnyEFz61xbaPudfsXbKYUphOg-qMDHO8g,33203
 nshtrainer/trainer/_distributed_prediction_result.py,sha256=bQw8Z6PT694UUf-zQPkech6CxyUSy8bAIexfSfPej0U,2507
 nshtrainer/trainer/_log_hparams.py,sha256=XH2lZ4U_3AZBhOt91ocsEhdL_NRz35oWvqLCUFDohUs,2389
@@ -160,6 +160,6 @@ nshtrainer/util/seed.py,sha256=diMV8iwBKN7Xxt5pELmui-gyqyT80_CZzomrWhNss0k,316
 nshtrainer/util/slurm.py,sha256=HflkP5iI_r4UHMyPjw9R4dD5AHsJUpcfJw5PLvGYBRM,1603
 nshtrainer/util/typed.py,sha256=Xt5fUU6zwLKSTLUdenovnKK0N8qUq89Kddz2_XeykVQ,164
 nshtrainer/util/typing_utils.py,sha256=MjY-CUX9R5Tzat-BlFnQjwl1PQ_W2yZQoXhkYHlJ_VA,442
-nshtrainer-1.3.0.dist-info/METADATA,sha256=M84AwXCuoJp21_m2IQKYDC-SFWDAdhOy-2fDL1jk9Lw,960
-nshtrainer-1.3.0.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-nshtrainer-1.3.0.dist-info/RECORD,,
+nshtrainer-1.3.2.dist-info/METADATA,sha256=XQdG9IP0N83areh70D1kM_rneikv2JQnO2VS34_MXRM,960
+nshtrainer-1.3.2.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+nshtrainer-1.3.2.dist-info/RECORD,,