suggestions

Author: vyron-arvanitis

bde/bde.py

@@ -239,10 +239,10 @@ def _draw_samples(
         self,
         logpost: Callable,
         rng_keys_e: ArrayLike,
-        init_positions_e,
-        L_e,
-        step_e,
-        sqrt_diag_e,
+        init_positions_e: ArrayLike,
+        L_e: ArrayLike,
+        step_e: ArrayLike,
+        sqrt_diag_e: ArrayLike,
     ):
         """Generate posterior samples for each ensemble member.
 
@@ -254,11 +254,11 @@ def _draw_samples(
             PRNG keys for each chain.
         init_positions_e : ParamTree
             Warmed-up starting positions with leading ensemble axis.
-        L_e : jax.Array
+        L_e : ArrayLike
             Inverse mass matrix factors per chain.
-        step_e : jax.Array
+        step_e : ArrayLike
             Step sizes per chain.
-        sqrt_diag_e : jax.Array
+        sqrt_diag_e : ArrayLike
             Square root of diagonal covariance estimates.
 
         Returns

bde/training/trainer.py

@@ -5,6 +5,8 @@
 from bde.loss.loss import BaseLoss, CategoricalCrossEntropy, GaussianNLL
 from bde.task import TaskType
 
+class BaseTrainer(): #TODO: maybe later?
+    pass
 
 class FnnTrainer:
     def __init__(self):