From 172c941cabefab0b36f32f3f9e5c13577e26491c Mon Sep 17 00:00:00 2001
From: Jackie <jackie.blaum@gmail.com>
Date: Tue, 5 Aug 2025 12:54:59 -0700
Subject: [PATCH 01/11] Fix: Handle zero valid samples in importance_reweight
 to avoid ValueError crash

---
 src/nbi/engine.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/nbi/engine.py b/src/nbi/engine.py
index 0faeb9b..5e58a5e 100644
--- a/src/nbi/engine.py
+++ b/src/nbi/engine.py
@@ -633,6 +633,10 @@ def importance_reweight(self, x_obs, x, y):
         log_weights = loglike + logprior - logproposal
 
         bad = np.isnan(log_weights) + np.isinf(log_weights)
+        valid_weights = log_weights[~bad]
+        if len(valid_weights) == 0:
+            print("All log weights are NaN or Inf — skipping this round!")
+            return np.zeros_like(log_weights)
         log_weights -= log_weights[~bad].max()
 
         weights = np.exp(log_weights)

From 79ea3935bd571212e9c0920be785246b31abf9aa Mon Sep 17 00:00:00 2001
From: Jackie <jackie.blaum@gmail.com>
Date: Wed, 13 Aug 2025 14:04:04 -0700
Subject: [PATCH 02/11] Added checks to ensure x_mean, x_std, y_mean, and y_std
 are set before scaling.

---
 src/nbi/engine.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/nbi/engine.py b/src/nbi/engine.py
index 5e58a5e..cc2ec6a 100644
--- a/src/nbi/engine.py
+++ b/src/nbi/engine.py
@@ -802,6 +802,7 @@ def scale_y(self, y, back=False):
         -------
 
         """
+        assert self.y_mean is not None and self.y_std is not None
         if back:
             return y * self.y_std + self.y_mean
         else:
@@ -824,6 +825,7 @@ def scale_x(self, x, back=False):
         -------
 
         """
+        assert self.x_mean is not None and self.x_std is not None
         if back:
             return x * self.x_std + self.x_mean
         else:

From 7193851a1b96e81c7ba579e36c40795e131c0bec Mon Sep 17 00:00:00 2001
From: Jackie <jackie.blaum@gmail.com>
Date: Wed, 13 Aug 2025 14:09:44 -0700
Subject: [PATCH 03/11] Added NaN guards to _train_step.

---
 src/nbi/engine.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/src/nbi/engine.py b/src/nbi/engine.py
index cc2ec6a..b0f505b 100644
--- a/src/nbi/engine.py
+++ b/src/nbi/engine.py
@@ -1038,10 +1038,32 @@ def _train_step(self):
         with self.tqdm(total=len(self.train_loader.dataset)) as pbar:
             for batch_idx, data in enumerate(self.train_loader):
                 x, y = data
+
+                if not torch.isfinite(x).all(): 
+                    raise ValueError(f"NaN/Inf in raw x @batch {batch_idx}")
+                if not torch.isfinite(y).all(): 
+                    raise ValueError(f"NaN/Inf in raw y @batch {batch_idx}")
+                
                 x = self.scale_x(x).to(self.device, dtype=torch.float32)
                 y = self.scale_y(y).to(self.device, dtype=torch.float32)
+                
+                if not torch.isfinite(x).all():
+                    bad_dims = (~torch.isfinite(x)).any(dim=0).nonzero(as_tuple=True)[0]
+                    raise ValueError(f"Scaling produced NaN/Inf in x @ dims {bad_dims.tolist()}")
+                if not torch.isfinite(y).all():
+                    bad_dims = (~torch.isfinite(y)).any(dim=0).nonzero(as_tuple=True)[0]
+                    raise ValueError(f"Scaling produced NaN/Inf in y @ dims {bad_dims.tolist()}")
+
                 self.optimizer.zero_grad()
                 loss = self.network(x, y)
+
+                if not torch.isfinite(loss).all():
+                    with torch.no_grad():
+                        print("Loss contains non-finite values. Sample stats:",
+                            "x mean/std", x.mean().item(), x.std().item(),
+                            "y mean/std", y.mean().item(), y.std().item())
+                    raise ValueError(f"Non-finite loss @batch {batch_idx}")
+    
                 loss = loss.mean()
                 train_loss.append(loss.item())
                 loss.backward()

From 5534981bcbfe3ba32bef7b2db1ea92745344779a Mon Sep 17 00:00:00 2001
From: Jackie <jackie.blaum@gmail.com>
Date: Sat, 21 Feb 2026 21:47:03 -0800
Subject: [PATCH 04/11] Add DatasetContainer to support torch Dataset
 integration

Introduces DatasetContainer class that wraps torch Datasets and provides
the same get_splits() API as BaseContainer, enabling NBI to reuse
_init_loader() with external torch datasets.
---
 src/nbi/data.py | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/src/nbi/data.py b/src/nbi/data.py
index 9048145..00ff622 100644
--- a/src/nbi/data.py
+++ b/src/nbi/data.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 from torch.utils.data import Dataset
+from torch.utils.data import Subset
 
 
 class Data:
@@ -55,3 +56,33 @@ def __getitem__(self, i, **kwargs):
             x, y = self.process(x, y)
 
         return np.atleast_2d(x), y
+
+class DatasetContainer:
+    """
+    Minimal container wrapper for torch Datasets so NBI can reuse _init_loader().
+    Provides the same get_splits() API as BaseContainer.
+    """
+    def __init__(self, dataset, f_val=0.1, f_test=0.0, seed=0):
+        self.dataset = dataset
+        self.f_val = float(f_val)
+        self.f_test = float(f_test)
+        self.seed = int(seed)
+
+    def get_splits(self):
+        n = len(self.dataset)
+        rng = np.random.default_rng(self.seed)
+        idx = np.arange(n)
+        rng.shuffle(idx)
+
+        n_test = int(self.f_test * n)
+        n_val  = int(self.f_val * n)
+
+        test_idx  = idx[:n_test]
+        val_idx   = idx[n_test:n_test + n_val]
+        train_idx = idx[n_test + n_val:]
+
+        return (
+            Subset(self.dataset, train_idx),
+            Subset(self.dataset, val_idx),
+            Subset(self.dataset, test_idx),
+        )
\ No newline at end of file

From c20942ca02d490fc34ae866b59612d8d7740e4e1 Mon Sep 17 00:00:00 2001
From: Jackie <jackie.blaum@gmail.com>
Date: Sat, 21 Feb 2026 21:55:34 -0800
Subject: [PATCH 05/11] Add torch Dataset and EmpiricalPrior support to NBI
 engine

- Support torch Datasets in _train_round() using DatasetContainer
- Integrate EmpiricalPrior handling in get_round_data() and log_prior()
- Add _check_scalers() validation to catch non-finite or zero scalers
- Add debug print statements for troubleshooting
---
 src/nbi/engine.py | 40 ++++++++++++++++++++++++++++++++++------
 1 file changed, 34 insertions(+), 6 deletions(-)

diff --git a/src/nbi/engine.py b/src/nbi/engine.py
index b0f505b..fc77a7a 100644
--- a/src/nbi/engine.py
+++ b/src/nbi/engine.py
@@ -21,6 +21,7 @@
 
 # this seems to be required for some environments
 from torch.utils.data import DataLoader, dataloader
+from torch.utils.data import Dataset as TorchDataset
 from tqdm import tqdm
 from tqdm.notebook import tqdm as tqdmn
 
@@ -28,6 +29,7 @@
 
 from .data import BaseContainer
 from .model import DataParallelFlow, get_featurizer, get_flow
+from .empirical_prior import EmpiricalPrior
 from .utils import iid_gaussian, log_like_iidg, parallel_simulate
 
 
@@ -388,9 +390,14 @@ def fit(
             self._init_train(lr)
             self._init_scheduler(min_lr, decay_type=decay_type)
             x_round, y_round = self.get_round_data(n_reuse)
-            data_container = BaseContainer(
-                x_round, y_round, f_test=0, f_val=f_val, process=self.process
-            )
+
+            if isinstance(x_round, TorchDataset) and y_round is None:
+                data_container = DatasetContainer(x_round, f_test=0.0, f_val=f_val, seed=0)
+            else:
+                data_container = BaseContainer(
+                    x_round, y_round, f_test=0, f_val=f_val, process=self.process
+                )
+
             self._init_loader(data_container, batch_size, workers=workers)
 
             for epoch in range(n_epochs):
@@ -900,6 +907,7 @@ def predict(
         x_path, good = self.simulate(ys)
         x_path = x_path[good]
         ys = ys[good]
+        print(ys)
         weights = self.importance_reweight(x, x_path, ys)
 
         neff = 1 / (weights**2).sum() - 1
@@ -1000,6 +1008,7 @@ def simulate(self, thetas):
             return self.x, masks
         else:
             n = len(thetas)
+            print(thetas)
             paths = np.array(
                 [os.path.join(path_round, str(i) + ".npy") for i in range(n)]
             )
@@ -1023,6 +1032,17 @@ def simulate(self, thetas):
                 masks = p.map(parallel_simulate, jobs)
             masks = np.concatenate(masks)
             return paths, masks
+        
+    def _check_scalers(self):
+        bad_mean = (~np.isfinite(self.x_mean)).any()
+        bad_std  = (~np.isfinite(self.x_std)).any()
+        if bad_mean or bad_std:
+            print("x_mean bad cols:", np.where(~np.isfinite(self.x_mean))[0].tolist())
+            print("x_std  bad cols:", np.where(~np.isfinite(self.x_std))[0].tolist())
+            raise ValueError("Non-finite x_mean/x_std; recompute scalers before training.")
+        if (self.x_std == 0).any():
+            zeros = np.nonzero(self.x_std == 0, as_tuple=False).T[-1].tolist()
+            print("x_std zero cols:", zeros)   
 
     def _train_step(self):
         """
@@ -1032,6 +1052,7 @@ def _train_step(self):
         -------
 
         """
+        self._check_scalers()
         np.random.seed(self.epoch)
         self.network.train()
         train_loss = []
@@ -1270,9 +1291,12 @@ def _draw_params(self, x, n):
                 return self.y
             else:
                 params = []
-                for prior in self.prior:
-                    params.append(prior.rvs(n))
-                params = np.array(params).T
+                if isinstance(self.prior, EmpiricalPrior):
+                    params = self.prior.rvs(n)
+                else:
+                    for prior in self.prior:
+                        params.append(prior.rvs(n))
+                    params = np.array(params).T
                 return params
         # 2+ round: sample from surrogate posterior
         else:
@@ -1280,7 +1304,9 @@ def _draw_params(self, x, n):
             logprior = self.log_prior(params)
             if np.isinf(logprior).any():
                 print("Samples outside prior N =", np.isinf(logprior).sum())
+                print('Bad: ', params[np.isinf(logprior)])
                 params = params[~np.isinf(logprior)]
+                print('Good: ', params)
                 while len(params) < n:
                     n_needed = n - len(params)
                     new_params = self.sample(x, n=n)
@@ -1332,6 +1358,8 @@ def log_prior(self, y):
         """
         if self.prior is None:
             return np.zeros(len(y))
+        elif isinstance(self.prior, EmpiricalPrior):
+            log_prob = self.prior.logpdf(y)
         else:
             log_prob = np.zeros(len(y))
             for i, prior in enumerate(self.prior):

From 133aac2107db4c85b954836174fd8d545c0f1f6c Mon Sep 17 00:00:00 2001
From: Jackie <jackie.blaum@gmail.com>
Date: Mon, 23 Feb 2026 13:50:09 -0800
Subject: [PATCH 06/11] Add multi-featurizer and multi-channel x support

---
 src/nbi/empirical_prior.py |  24 ++++
 src/nbi/engine.py          | 235 +++++++++++++++++++++++++++++++------
 src/nbi/model.py           |  48 +++++++-
 src/nbi/utils.py           |  68 +++++++++++
 4 files changed, 340 insertions(+), 35 deletions(-)
 create mode 100644 src/nbi/empirical_prior.py

diff --git a/src/nbi/empirical_prior.py b/src/nbi/empirical_prior.py
new file mode 100644
index 0000000..ec22946
--- /dev/null
+++ b/src/nbi/empirical_prior.py
@@ -0,0 +1,24 @@
+import numpy as np
+
+class EmpiricalPrior:
+    def __init__(self, lookup_table, priors=None, random_state=None):
+        self.lookup_table = np.asarray(lookup_table)
+        self.n_samples = self.lookup_table.shape[0]
+        self.priors = priors
+        self.rng = np.random.default_rng(random_state)
+
+    def rvs(self, size=1, replace=True):
+        if not replace and size > self.n_samples:
+            raise ValueError(
+                f"Requested {size} samples without replacement, but "
+                f"lookup_table only has {self.n_samples} entries."
+            )
+        idx = self.rng.choice(self.n_samples, size=size, replace=replace)
+        return self.lookup_table[idx]
+
+    def logpdf(self, params):
+        if self.priors is None:
+            raise NotImplementedError(
+                "logp not implemented: provide a prior model or override this method."
+            )
+        return self.priors.logpdf(params)
diff --git a/src/nbi/engine.py b/src/nbi/engine.py
index fc77a7a..b4dbf3c 100644
--- a/src/nbi/engine.py
+++ b/src/nbi/engine.py
@@ -27,11 +27,10 @@
 
 dataloader.multiprocessing = mp
 
-from .data import BaseContainer
+from .data import BaseContainer, DatasetContainer
 from .model import DataParallelFlow, get_featurizer, get_flow
 from .empirical_prior import EmpiricalPrior
-from .utils import iid_gaussian, log_like_iidg, parallel_simulate
-
+from .utils import iid_gaussian, log_like_iidg, parallel_simulate, collate_list_channels, masked_mean_std
 
 class NBI:
     """Neural Bayesian Inference Engine.
@@ -143,11 +142,13 @@ def __init__(
         path="test",
         n_jobs=1,
         labels=None,
+        masked_channels=None,
         tqdm_notebook=False,
         network_reinit=False,
         scale_reinit=True,
     ):
         self.device = device
+        self.masked_channels = masked_channels
         self.init_env()
 
         if state_dict is not None:
@@ -164,7 +165,29 @@ def __init__(
             flow_config_all = copy.copy(self.default_flow_config)
             flow_config_all.update(flow)
 
-        if type(featurizer) == dict:
+        featurizer_config = copy.deepcopy(featurizer)
+
+        if isinstance(featurizer, list):
+            built = []
+            for i, f in enumerate(featurizer):
+                if isinstance(f, dict):
+                    if "type" not in f:
+                        raise ValueError(f"Featurizer config {i} missing 'type': {f}")
+                    built.append(get_featurizer(f["type"], f))
+                elif isinstance(f, nn.Module):
+                    built.append(f)
+                else:
+                    raise TypeError(f"Featurizer {i} must be dict or nn.Module, got {type(f)}")
+            featurizer = built
+
+            # strict: no Nones
+            if any(f is None for f in featurizer):
+                bad = [i for i, f in enumerate(featurizer) if f is None]
+                raise ValueError(f"Some featurizers built to None at indices {bad}")
+
+            print("Built featurizers:", len(featurizer))
+
+        elif isinstance(featurizer, dict):
             featurizer = get_featurizer(featurizer["type"], featurizer)
             flow_config_all["num_cond_inputs"] = featurizer.num_outputs
 
@@ -172,9 +195,7 @@ def __init__(
         self.flow_config = flow_config_all
 
         self.network = get_flow(featurizer, **flow_config_all)
-        self.network = DataParallelFlow(self.network).to(
-            self.device, dtype=torch.float32
-        )
+        self.network = DataParallelFlow(self.network).to(self.device, dtype=torch.float32)
 
         self.corner_kwargs.update({"labels": labels})
         self.network_reinit = network_reinit
@@ -367,6 +388,8 @@ def fit(
         self.x = x
         self.y = y
 
+        using_dataset = isinstance(x, TorchDataset) and y is None
+
         # this needs revision in another version
         self.wandb = use_wandb
         if self.wandb:
@@ -376,8 +399,10 @@ def fit(
             min_lr = min(lr, lr / (n_sims / batch_size * n_epochs) * 10)
             print("Auto learning rate to min_lr =", min_lr)
 
+        using_dataset = isinstance(self.x, TorchDataset) and self.y is None
+
         # for restarting training
-        if len(self.x_all) == self.round:
+        if (not using_dataset) and (len(self.x_all) == self.round):
             # this is not a restart because
             # data for this round has not been generated
             self.prepare_data(x_obs, n_sims)
@@ -600,6 +625,10 @@ def get_round_data(self, n_reuse):
             Training parameters for the current round.
 
         """
+        # if user provided a torch Dataset for amortized training
+        if isinstance(self.x, TorchDataset) and self.y is None and self.round == 0:
+            return self.x, None
+
         if n_reuse == -1:
             return np.concatenate(self.x_all), np.concatenate(self.y_all)
         else:
@@ -821,6 +850,10 @@ def scale_x(self, x, back=False):
         """
         Scale data to zero mean and unit variance, and vice versa.
 
+        Supports:
+        - x as Tensor/ndarray  [B, C, L] or [C, L]
+        - x as list/tuple of Tensors/arrays, each [B, C, L] or [C, L]
+
         Parameters
         ----------
         x : ndarray
@@ -833,6 +866,28 @@ def scale_x(self, x, back=False):
 
         """
         assert self.x_mean is not None and self.x_std is not None
+
+        # Multi-channel list case
+        if isinstance(x, (list, tuple)):
+            assert isinstance(self.x_mean, list) and isinstance(self.x_std, list)
+            assert len(x) == len(self.x_mean) == len(self.x_std)
+
+            xs = []
+            for j, xj in enumerate(x):
+                mu = self.x_mean[j]
+                sd = self.x_std[j]
+
+                # If xj is torch Tensor, keep it torch
+                if torch.is_tensor(xj):
+                    mu_t = torch.as_tensor(mu, device=xj.device, dtype=xj.dtype)
+                    sd_t = torch.as_tensor(sd, device=xj.device, dtype=xj.dtype)
+                    xs.append(xj * sd_t + mu_t if back else (xj - mu_t) / sd_t)
+                else:
+                    xs.append(xj * sd + mu if back else (xj - mu) / sd)
+
+            return xs
+
+        # Single tensor/array case (original behavior)
         if back:
             return x * self.x_std + self.x_mean
         else:
@@ -1032,8 +1087,38 @@ def simulate(self, thetas):
                 masks = p.map(parallel_simulate, jobs)
             masks = np.concatenate(masks)
             return paths, masks
+
         
     def _check_scalers(self):
+        if self.x_mean is None or self.x_std is None:
+            raise ValueError("x_mean/x_std are None; did _init_scales run?")
+
+        # Multi-channel case: list of [1, C, 1] arrays
+        if isinstance(self.x_mean, list):
+            if not isinstance(self.x_std, list):
+                raise ValueError("x_mean is a list but x_std is not.")
+            if len(self.x_mean) != len(self.x_std):
+                raise ValueError("x_mean/x_std length mismatch.")
+
+            for j, (mu, sd) in enumerate(zip(self.x_mean, self.x_std)):
+                mu = np.asarray(mu)
+                sd = np.asarray(sd)
+
+                if (~np.isfinite(mu)).any():
+                    bad = np.where(~np.isfinite(mu))
+                    raise ValueError(f"Non-finite x_mean in channel {j} at {bad}")
+
+                if (~np.isfinite(sd)).any():
+                    bad = np.where(~np.isfinite(sd))
+                    raise ValueError(f"Non-finite x_std in channel {j} at {bad}")
+
+                if (sd == 0).any():
+                    bad = np.where(sd == 0)
+                    raise ValueError(f"Zero x_std in channel {j} at {bad}")
+
+            return  # all good
+
+        # Single-tensor case (original behavior)
         bad_mean = (~np.isfinite(self.x_mean)).any()
         bad_std  = (~np.isfinite(self.x_std)).any()
         if bad_mean or bad_std:
@@ -1041,8 +1126,9 @@ def _check_scalers(self):
             print("x_std  bad cols:", np.where(~np.isfinite(self.x_std))[0].tolist())
             raise ValueError("Non-finite x_mean/x_std; recompute scalers before training.")
         if (self.x_std == 0).any():
-            zeros = np.nonzero(self.x_std == 0, as_tuple=False).T[-1].tolist()
-            print("x_std zero cols:", zeros)   
+            zeros = np.where(self.x_std == 0)[0].tolist()
+            print("x_std zero cols:", zeros)
+            raise ValueError("Zero x_std; recompute scalers before training.") 
 
     def _train_step(self):
         """
@@ -1060,17 +1146,35 @@ def _train_step(self):
             for batch_idx, data in enumerate(self.train_loader):
                 x, y = data
 
-                if not torch.isfinite(x).all(): 
-                    raise ValueError(f"NaN/Inf in raw x @batch {batch_idx}")
+                # Raw x finite check (supports list)
+                if isinstance(x, (list, tuple)):
+                    for j, xj in enumerate(x):
+                        if not torch.isfinite(xj).all():
+                            raise ValueError(f"NaN/Inf in raw x channel {j} @batch {batch_idx}")
+                else:
+                    if not torch.isfinite(x).all():
+                        raise ValueError(f"NaN/Inf in raw x @batch {batch_idx}")
                 if not torch.isfinite(y).all(): 
                     raise ValueError(f"NaN/Inf in raw y @batch {batch_idx}")
                 
-                x = self.scale_x(x).to(self.device, dtype=torch.float32)
+                # scale + move x (supports list-of-tensors)
+                x = self.scale_x(x)
+                if isinstance(x, (list, tuple)):
+                    x = [xj.to(self.device, dtype=torch.float32) for xj in x]
+                else:
+                    x = x.to(self.device, dtype=torch.float32)
+
+                # scale + move y
                 y = self.scale_y(y).to(self.device, dtype=torch.float32)
                 
-                if not torch.isfinite(x).all():
-                    bad_dims = (~torch.isfinite(x)).any(dim=0).nonzero(as_tuple=True)[0]
-                    raise ValueError(f"Scaling produced NaN/Inf in x @ dims {bad_dims.tolist()}")
+                if isinstance(x, (list, tuple)):
+                    for j, xj in enumerate(x):
+                        if not torch.isfinite(xj).all():
+                            raise ValueError(f"Scaling produced NaN/Inf in x channel {j} @batch {batch_idx}")
+                else:
+                    if not torch.isfinite(x).all():
+                        bad_dims = (~torch.isfinite(x)).any(dim=0).nonzero(as_tuple=True)[0]
+                        raise ValueError(f"Scaling produced NaN/Inf in x @ dims {bad_dims.tolist()}")
                 if not torch.isfinite(y).all():
                     bad_dims = (~torch.isfinite(y)).any(dim=0).nonzero(as_tuple=True)[0]
                     raise ValueError(f"Scaling produced NaN/Inf in y @ dims {bad_dims.tolist()}")
@@ -1096,7 +1200,12 @@ def _train_step(self):
                     )
                 self.optimizer.step()
 
-                pbar.update(x.shape[0])
+                # x can be a list of channel tensors
+                if isinstance(x, (list, tuple)):
+                    bs = x[0].shape[0]
+                else:
+                    bs = x.shape[0]
+                pbar.update(bs)
                 pbar.set_description(
                     "Epoch {:d}: Train, Loglike in nats: {:.6f}".format(
                         self.epoch, -np.mean(train_loss)
@@ -1123,14 +1232,23 @@ def _validate_step(self):
             objs = 0
             for batch_idx, data in enumerate(self.valid_loader):
                 x, y = data
-                x = self.scale_x(x).to(self.device, dtype=torch.float32)
+                x = self.scale_x(x)
+                if isinstance(x, (list, tuple)):
+                    x = [xj.to(self.device, dtype=torch.float32) for xj in x]
+                else:
+                    x = x.to(self.device, dtype=torch.float32)
+
                 y = self.scale_y(y).to(self.device, dtype=torch.float32)
-                objs += x.shape[0]
+                if isinstance(x, (list, tuple)):
+                    bs = x[0].shape[0]
+                else:
+                    bs = x.shape[0]
+                objs += bs
                 self.optimizer.zero_grad()
                 with torch.no_grad():
                     loss = self.network(x, y).mean()
                     val_loss.append(loss.detach().cpu().numpy())
-                pbar.update(x.shape[0])
+                pbar.update(bs)
                 pbar.set_description(
                     f"- Val, Loglike in nats: {-np.sum(val_loss) / (batch_idx + 1):.6f}"
                 )
@@ -1259,6 +1377,7 @@ def _init_loader(self, data_container, batch_size, workers=4):
             "pin_memory": False,
             "drop_last": True,
             "persistent_workers": True,
+            "collate_fn": collate_list_channels, 
         }
 
         self.train_loader = DataLoader(
@@ -1318,28 +1437,76 @@ def _draw_params(self, x, n):
 
     def _init_scales(self):
         """
-            Calculate data pre-processing scales from the current round training data.
-
-        Returns
-        -------
-
+        Calculate data pre-processing scales from the current round training data.
+        Supports x being either:
+        - a Tensor [B, C, L]
+        - a list of Tensors, each [B, C, L] (multi-featurizer / multi-channel)
         """
-        x_list = []
+        x_batches = []
         y_list = []
         n = 0
+
         for batch_idx, data in enumerate(self.train_loader):
             x, y = data
-            x_list.append(x.cpu().numpy())
             y_list.append(y.cpu().numpy())
-            n += x_list[-1].shape[0]
+            x_batches.append(x)
+            n += y.shape[0]
             if n > 5000:
                 break
-        x_list = np.concatenate(x_list, axis=0)
-        y_list = np.concatenate(y_list, axis=0)
-        self.x_mean = x_list.mean(-1, keepdims=True).mean(0, keepdims=True)
-        self.x_std = x_list.std(-1, keepdims=True).mean(0, keepdims=True)
-        self.y_mean = y_list.mean(0, keepdims=True)
-        self.y_std = y_list.std(0, keepdims=True)
+
+        # y scales (same as before)
+        y_arr = np.concatenate(y_list, axis=0)
+        self.y_mean = y_arr.mean(0, keepdims=True)
+        self.y_std  = y_arr.std(0, keepdims=True)
+
+        # x scales
+        x0 = x_batches[0]
+
+        # Case 1: multi-channel list/tuple of tensors
+        if isinstance(x0, (list, tuple)):
+            x_mean = []
+            x_std = []
+            n_chan = len(x0)
+
+            for j in range(n_chan):
+                # concatenate batches along batch dim: [N, C, L]
+                xs = torch.cat([xb[j] for xb in x_batches], dim=0).cpu().numpy()
+
+                cfg = None
+                if self.masked_channels is not None:
+                    cfg = self.masked_channels[j]
+
+                if cfg is not None:
+                    mu, sd = masked_mean_std(
+                        xs,
+                        mask_dim=cfg["mask_dim"],
+                        masked_dims=cfg["masked_dims"],
+                    )
+
+                    if cfg.get("mask_keep_identity", True):
+                        md = cfg["mask_dim"]
+                        mu[0, md, 0] = 0.0
+                        sd[0, md, 0] = 1.0
+
+                else:
+                    # default unmasked behavior
+                    mu = xs.mean(axis=(0, 2), keepdims=True)
+                    sd = xs.std(axis=(0, 2), keepdims=True)
+                    sd = np.where(sd == 0, 1.0, sd)
+
+                x_mean.append(mu.astype(np.float32))
+                x_std.append(sd.astype(np.float32))
+
+            self.x_mean = x_mean
+            self.x_std = x_std
+
+        # Case 2: single tensor [B, C, L]
+        else:
+            x_arr = torch.cat([xb for xb in x_batches], dim=0).cpu().numpy()
+            self.x_mean = x_arr.mean(-1, keepdims=True).mean(0, keepdims=True)
+            self.x_std  = x_arr.std(-1, keepdims=True).mean(0, keepdims=True)
+            self.x_std  = np.where(self.x_std == 0, 1.0, self.x_std).astype(np.float32)
+            self.x_mean = self.x_mean.astype(np.float32)
 
     def log_prior(self, y):
         """
diff --git a/src/nbi/model.py b/src/nbi/model.py
index be8c10f..9081f31 100644
--- a/src/nbi/model.py
+++ b/src/nbi/model.py
@@ -1,7 +1,36 @@
 from torch import nn
+import torch
 
 from .nn import RNN, ResNet, flows
 
+class MultiFeaturizer(nn.Module):
+    """
+    Wraps a list of featurizers into a single module.
+
+    Expects x to be a list/tuple of tensors, one per channel.
+    Returns concatenated feature vector: [B, sum_j out_j]
+    """
+    def __init__(self, featurizers):
+        super().__init__()
+        self.featurizers = nn.ModuleList(featurizers)
+        # expose num_outputs so flow can size conditioning dim
+        self.num_outputs = int(sum(getattr(f, "num_outputs", 0) for f in featurizers))
+
+    def forward(self, x):
+        if not isinstance(x, (list, tuple)):
+            raise TypeError(f"MultiFeaturizer expects x as list/tuple, got {type(x)}")
+
+        if len(x) != len(self.featurizers):
+            raise ValueError(f"x has {len(x)} channels but featurizers has {len(self.featurizers)}")
+
+        feats = []
+        for fj, xj in zip(self.featurizers, x):
+            out = fj(xj)
+            # flatten to [B, -1] if needed
+            if out.ndim > 2:
+                out = out.reshape(out.shape[0], -1)
+            feats.append(out)
+        return torch.cat(feats, dim=1)
 
 class DataParallelFlow(nn.DataParallel):
     def __init__(self, *args, **kwargs):
@@ -85,19 +114,36 @@ def get_featurizer(network_type, config):
             bidirectional=False,
             rnn="GRU",
         )
+    else:
+        raise ValueError(f"Unknown featurizer type: {network_type}")
 
 
 def get_flow(
     featurizer,
     n_dims,
     flow_hidden,
-    num_cond_inputs,
+    num_cond_inputs=None,
     num_blocks=5,
     perm_seed=0,
     clamp_0=-1,
     clamp_1=1,
     n_mog=8,
 ):
+
+    # If a list/tuple of per-channel featurizers is provided, wrap them
+    if isinstance(featurizer, (list, tuple)):
+        featurizer = MultiFeaturizer(featurizer)
+
+    if num_cond_inputs is None:
+        if featurizer is None:
+            raise ValueError("num_cond_inputs must be provided when featurizer is None")
+        if not hasattr(featurizer, "num_outputs"):
+            raise ValueError(
+                "Could not infer num_cond_inputs: featurizer has no attribute 'num_outputs'. "
+                "Pass num_cond_inputs explicitly."
+            )
+        num_cond_inputs = int(featurizer.num_outputs)
+
     modules = []
     MADE = flows.MADE2
     num_blocks -= 1
diff --git a/src/nbi/utils.py b/src/nbi/utils.py
index c1c610b..be19551 100644
--- a/src/nbi/utils.py
+++ b/src/nbi/utils.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 from tqdm import tqdm
+import torch
 
 
 def parallel_simulate(args):
@@ -51,3 +52,70 @@ def log_like(x_err, x, x_path, y):
 
 def log_like_iidg(x_err):
     return partial(log_like, x_err)
+
+def collate_list_channels(batch):
+    """
+    Batch is a list of tuples: [(x_list, y), ...]
+    where x_list is list of tensors: [x0, x1, ..., xK]
+    Returns:
+    x_batched: list of tensors, each stacked to [B, ...]
+    y_batched: tensor [B, D]
+    """
+    xs, ys = zip(*batch)  # xs is tuple of lists
+    k = len(xs[0])
+    x_out = [torch.stack([x_i[j] for x_i in xs], dim=0) for j in range(k)]
+    y_out = torch.stack(list(ys), dim=0)
+    return x_out, y_out
+
+
+def masked_mean_std(xs, mask_dim, masked_dims):
+    """
+    xs: np.ndarray [N, C, L]
+    mask_dim: int
+    masked_dims: list[int]
+    Returns mu, sd shaped [1, C, 1]
+    """
+    C = xs.shape[1]
+    mu = np.zeros((1, C, 1), dtype=np.float32)
+    sd = np.ones((1, C, 1), dtype=np.float32)
+
+    mask = xs[:, mask_dim, :]  # [N, L]
+    # force mask to {0,1} and finite
+    mask = np.where(np.isfinite(mask), mask, 0.0)
+    mask = (mask > 0.5).astype(np.float32)
+
+    for d in masked_dims:
+        v = xs[:, d, :]
+        # treat non-finite values as missing
+        finite = np.isfinite(v)
+        w = mask * finite.astype(np.float32)
+        wsum = w.sum()
+
+        if wsum <= 0:
+            # fallback: unmasked finite-only stats
+            vv = v[finite]
+            if vv.size == 0:
+                mu[0, d, 0] = 0.0
+                sd[0, d, 0] = 1.0
+                continue
+            m = float(vv.mean())
+            s = float(vv.std())
+            if (not np.isfinite(s)) or s == 0:
+                s = 1.0
+            mu[0, d, 0] = m
+            sd[0, d, 0] = s
+            continue
+
+        m = (v * w).sum() / wsum
+        var = ((v - m) ** 2 * w).sum() / wsum
+        s = np.sqrt(var)
+
+        if not np.isfinite(m):
+            m = 0.0
+        if not np.isfinite(s) or s == 0:
+            s = 1.0
+
+        mu[0, d, 0] = float(m)
+        sd[0, d, 0] = float(s)
+
+    return mu, sd
\ No newline at end of file

From 6eb12cbb765f202698a57ad6149ab786ac5de576 Mon Sep 17 00:00:00 2001
From: Jackie <jackie.blaum@gmail.com>
Date: Wed, 11 Mar 2026 12:39:52 -0700
Subject: [PATCH 07/11] Add documentation for new multi-channel and DataLoader
 parameters

- Document pin_memory and prefetch_factor in fit() and _init_loader()
- Update set_params() docstring to reflect new state dict format with legacy support
- Update _init_loader() data_container type hint
- Add TODO for generalizing add_noise to multi-channel data
---
 src/nbi/data.py   |  50 +++++++++++++++-
 src/nbi/engine.py | 149 +++++++++++++++++++++++++++++++---------------
 src/nbi/model.py  |   6 +-
 src/nbi/utils.py  | 116 +++++++++++++++++++++++-------------
 4 files changed, 228 insertions(+), 93 deletions(-)

diff --git a/src/nbi/data.py b/src/nbi/data.py
index 00ff622..6fc947d 100644
--- a/src/nbi/data.py
+++ b/src/nbi/data.py
@@ -57,16 +57,60 @@ def __getitem__(self, i, **kwargs):
 
         return np.atleast_2d(x), y
 
+class ProcessedTorchDataset(Dataset):
+    """
+    Wrap a torch Dataset and apply `process` on-the-fly in __getitem__.
+
+    Supports datasets that return:
+      - x
+      - (x, y)
+    and supports process signatures:
+      - process(x, y) -> (x2, y2)
+      - process(x)    -> x2
+    """
+    def __init__(self, dataset, process=None):
+        self.dataset = dataset
+        self.process = process
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        out = self.dataset[idx]
+
+        # Normalize to (x, y)
+        if isinstance(out, (tuple, list)) and len(out) == 2:
+            x, y = out
+        else:
+            x, y = out, None
+
+        if self.process is None:
+            return out
+
+        # Apply process
+        try:
+            outp = self.process(x, y)
+        except TypeError:
+            outp = self.process(x)
+
+        return outp
+
 class DatasetContainer:
     """
-    Minimal container wrapper for torch Datasets so NBI can reuse _init_loader().
-    Provides the same get_splits() API as BaseContainer.
+    Container wrapper for torch Datasets so NBI can reuse _init_loader().
+    Provides the same get_splits() API as BaseContainer, and supports `process`
+    applied in __getitem__ (like BaseContainer).
     """
-    def __init__(self, dataset, f_val=0.1, f_test=0.0, seed=0):
+    def __init__(self, dataset, f_val=0.1, f_test=0.0, seed=0, process=None):
         self.dataset = dataset
         self.f_val = float(f_val)
         self.f_test = float(f_test)
         self.seed = int(seed)
+        self.process = process
+
+        # Wrap so process is applied for any subset indexing
+        if self.process is not None:
+            self.dataset = ProcessedTorchDataset(self.dataset, self.process)
 
     def get_splits(self):
         n = len(self.dataset)
diff --git a/src/nbi/engine.py b/src/nbi/engine.py
index b4dbf3c..2d00803 100644
--- a/src/nbi/engine.py
+++ b/src/nbi/engine.py
@@ -194,8 +194,10 @@ def __init__(
         self.featurizer_config = copy.copy(featurizer)
         self.flow_config = flow_config_all
 
-        self.network = get_flow(featurizer, **flow_config_all)
-        self.network = DataParallelFlow(self.network).to(self.device, dtype=torch.float32)
+        self.network = get_flow(featurizer, **flow_config_all).to(self.device, dtype=torch.float32)
+
+        if "cuda" in self.device and torch.cuda.is_available() and torch.cuda.device_count() > 1:
+            self.network = DataParallelFlow(self.network)
 
         self.corner_kwargs.update({"labels": labels})
         self.network_reinit = network_reinit
@@ -279,6 +281,8 @@ def fit(
         plot=True,
         f_accept_min=-1,
         workers=8,
+        pin_memory=None,
+        prefetch_factor=2,
     ):
         """
         Fit the Neural Bayesian Inference Engine.
@@ -365,6 +369,12 @@ def fit(
         workers : int, optional
             Number of workers for data loading.
 
+        pin_memory : bool, optional
+            If True, enables pinned memory for data loading. Defaults to True when using CUDA, False otherwise.
+
+        prefetch_factor : int, optional
+            Number of batches prefetched per worker. Default: 2. Ignored when workers=0.
+
         Returns
         -------
 
@@ -408,6 +418,11 @@ def fit(
             self.prepare_data(x_obs, n_sims)
 
         for i in range(n_rounds):
+            path_round = os.path.join(self.directory, str(self.round))
+            try:
+                os.mkdir(path_round)
+            except:
+                pass
             print(
                 f"\n---------------------- Round: {self.round} ----------------------"
             )
@@ -417,13 +432,13 @@ def fit(
             x_round, y_round = self.get_round_data(n_reuse)
 
             if isinstance(x_round, TorchDataset) and y_round is None:
-                data_container = DatasetContainer(x_round, f_test=0.0, f_val=f_val, seed=0)
+                data_container = DatasetContainer(x_round, f_test=0.0, f_val=f_val, seed=0, process=self.process)
             else:
                 data_container = BaseContainer(
                     x_round, y_round, f_test=0, f_val=f_val, process=self.process
                 )
 
-            self._init_loader(data_container, batch_size, workers=workers)
+            self._init_loader(data_container, batch_size, workers=workers, pin_memory=pin_memory, prefetch_factor=prefetch_factor)
 
             for epoch in range(n_epochs):
                 self.epoch = epoch
@@ -759,56 +774,72 @@ def set_params(self, state_dict):
 
         Parameters
         ----------
-        state_dict : str or state dict
-            State dict or path to saved state dict containing three keys:
-            network_state_dict, x_scale, y_scale
+        state_dict : str or dict
+            State dict or path to saved state dict. Supports both the current format
+            (model_state_dict, x_mean, x_std, y_mean, y_std) and legacy format (x_scale, y_scale).
 
         Returns
         -------
 
         """
-
         if type(state_dict) == str:
-            state_dict = torch.load(
-                state_dict, map_location=self.device, weights_only=False
-            )
-        model_state_dict = state_dict["model_state_dict"]
-
-        # Move x_scale and y_scale to CPU before converting to numpy arrays
-        x_scale = state_dict["x_scale"].cpu().numpy()
-        y_scale = state_dict["y_scale"].cpu().numpy()
+            state_dict = torch.load(state_dict, map_location=self.device, weights_only=False)
+
+        self.get_network().load_state_dict(state_dict["model_state_dict"])
+
+        def to_numpy(obj):
+            if isinstance(obj, torch.Tensor):
+                return obj.detach().cpu().numpy()
+            if isinstance(obj, np.ndarray):
+                return obj
+            if isinstance(obj, (list, tuple)):
+                return [to_numpy(x) for x in obj]
+            if obj is None:
+                return None
+            raise TypeError(f"Unsupported scaler type: {type(obj)}")
+
+        if "x_mean" in state_dict:
+            self.x_mean = to_numpy(state_dict["x_mean"])
+            self.x_std  = to_numpy(state_dict["x_std"])
+            self.y_mean = to_numpy(state_dict["y_mean"])
+            self.y_std  = to_numpy(state_dict["y_std"])
+        else:
+            x_scale = state_dict["x_scale"].cpu().numpy()
+            y_scale = state_dict["y_scale"].cpu().numpy()
+            self.x_mean, self.x_std = x_scale[0], x_scale[1]
+            self.y_mean, self.y_std = y_scale[0], y_scale[1]
 
-        self.x_mean = x_scale[0]
-        self.x_std = x_scale[1]
-        self.y_mean = y_scale[0]
-        self.y_std = y_scale[1]
-        self.get_network().load_state_dict(model_state_dict)
+        # keep masks if present
+        if "masked_channels" in state_dict:
+            self.masked_channels = state_dict["masked_channels"]
 
     def get_params(self):
         """
         Saves the network weights and pre-processing scales to disk
-
-        Returns
-        -------
-
         """
-        x_scale = np.array([self.x_mean, self.x_std], dtype=np.float32)
-        y_scale = np.array([self.y_mean, self.y_std], dtype=np.float32)
-
-        # Convert numpy arrays to PyTorch tensors
-        x_scale_tensor = torch.from_numpy(x_scale)
-        y_scale_tensor = torch.from_numpy(y_scale)
-
-        # Assuming 'network' is your model
         model_state_dict = copy.deepcopy(self.get_network().state_dict())
 
-        # Create a new dictionary to store model state and additional tensors
+        def to_tensor(obj):
+            # obj can be ndarray, torch tensor, or list/tuple of those
+            if isinstance(obj, torch.Tensor):
+                return obj.detach().cpu()
+            if isinstance(obj, np.ndarray):
+                return torch.from_numpy(obj.astype(np.float32, copy=False))
+            if isinstance(obj, (list, tuple)):
+                return [to_tensor(x) for x in obj]
+            if obj is None:
+                return None
+            raise TypeError(f"Unsupported scaler type: {type(obj)}")
+
         state_dict = {
             "model_state_dict": model_state_dict,
-            "x_scale": x_scale_tensor,
-            "y_scale": y_scale_tensor,
+            "x_mean": to_tensor(self.x_mean),
+            "x_std":  to_tensor(self.x_std),
+            "y_mean": to_tensor(self.y_mean),
+            "y_std":  to_tensor(self.y_std),
             "flow_config": self.flow_config,
             "featurizer_config": self.featurizer_config,
+            "masked_channels": self.masked_channels,
         }
         return state_dict
 
@@ -1021,7 +1052,18 @@ def sample(self, x, y=None, n=5000, corner=False):
         """
         self.network.eval()
         x = self.scale_x(x)
-        x = torch.from_numpy(x).to(self.device, dtype=torch.float32)
+
+        if isinstance(x, (list, tuple)):
+            x = [
+                (xj if torch.is_tensor(xj) else torch.as_tensor(xj))
+                .to(self.device, dtype=torch.float32)
+                for xj in x
+            ]
+        else:
+            if not torch.is_tensor(x):
+                x = torch.as_tensor(x)
+            x = x.to(self.device, dtype=torch.float32)
+
         with torch.no_grad():
             # GPU memory control (make larger?)
             if n > 20000:
@@ -1353,18 +1395,22 @@ def _step_scheduler(self):
         else:
             self.scheduler.step()
 
-    def _init_loader(self, data_container, batch_size, workers=4):
+    def _init_loader(self, data_container, batch_size, workers=4, pin_memory=None, prefetch_factor=2):
         """
         Initialize data loader.
 
         Parameters
         ----------
-        data_container : DataContainer
+        data_container : BaseContainer or DatasetContainer
             Data container object.
         batch_size : int
             Batch size.
         workers : int, optional
             Number of workers for data loader.
+        pin_memory : bool, optional
+            If True, enables pinned memory for data loading. Defaults to True when using CUDA, False otherwise.
+        prefetch_factor : int, optional
+            Number of batches prefetched per worker. Default: 2. Ignored when workers=0.
 
         Returns
         -------
@@ -1372,13 +1418,20 @@ def _init_loader(self, data_container, batch_size, workers=4):
         """
         train_container, val_container, test_container = data_container.get_splits()
 
+        if pin_memory is None:
+            pin_memory = ("cuda" in str(self.device))
+
         kwargs = {
             "num_workers": workers,
-            "pin_memory": False,
+            "pin_memory": pin_memory,
             "drop_last": True,
             "persistent_workers": True,
+            "prefetch_factor": prefetch_factor if workers > 0 else None,
             "collate_fn": collate_list_channels, 
         }
+        if workers == 0:
+            kwargs.pop("prefetch_factor", None)
+            kwargs["persistent_workers"] = False
 
         self.train_loader = DataLoader(
             train_container, batch_size=batch_size, shuffle=True, **kwargs
@@ -1472,26 +1525,28 @@ def _init_scales(self):
                 # concatenate batches along batch dim: [N, C, L]
                 xs = torch.cat([xb[j] for xb in x_batches], dim=0).cpu().numpy()
 
-                cfg = None
-                if self.masked_channels is not None:
-                    cfg = self.masked_channels[j]
+                cfg = self.masked_channels[j] if self.masked_channels is not None else None
+                reduce_axes = (0, 2)
+                if cfg is not None and "reduce_axes" in cfg:
+                    reduce_axes = tuple(cfg["reduce_axes"])
 
                 if cfg is not None:
                     mu, sd = masked_mean_std(
                         xs,
                         mask_dim=cfg["mask_dim"],
                         masked_dims=cfg["masked_dims"],
+                        reduce_axes=reduce_axes
                     )
 
                     if cfg.get("mask_keep_identity", True):
                         md = cfg["mask_dim"]
-                        mu[0, md, 0] = 0.0
-                        sd[0, md, 0] = 1.0
+                        mu[0, md, :] = 0.0
+                        sd[0, md, :] = 1.0
 
                 else:
                     # default unmasked behavior
-                    mu = xs.mean(axis=(0, 2), keepdims=True)
-                    sd = xs.std(axis=(0, 2), keepdims=True)
+                    mu = xs.mean(axis=reduce_axes, keepdims=True)
+                    sd = xs.std(axis=reduce_axes, keepdims=True)
                     sd = np.where(sd == 0, 1.0, sd)
 
                 x_mean.append(mu.astype(np.float32))
diff --git a/src/nbi/model.py b/src/nbi/model.py
index 9081f31..9ca8dc9 100644
--- a/src/nbi/model.py
+++ b/src/nbi/model.py
@@ -33,8 +33,8 @@ def forward(self, x):
         return torch.cat(feats, dim=1)
 
 class DataParallelFlow(nn.DataParallel):
-    def __init__(self, *args, **kwargs):
-        super(type(self), self).__init__(*args, **kwargs)
+    def __init__(self, module, device_ids=None, output_device=None, dim=0):
+        super().__init__(module, device_ids=device_ids, output_device=output_device, dim=dim)
 
     def sample(self, x, n=1000, is_feature=False):
         if not is_feature:
@@ -46,7 +46,7 @@ def sample(self, x, n=1000, is_feature=False):
 
 class Flow(nn.Module):
     def __init__(self, featurizer, model):
-        super(type(self), self).__init__()
+        super().__init__()
         self.featurizer = featurizer
         self.flow = model
 
diff --git a/src/nbi/utils.py b/src/nbi/utils.py
index be19551..93b228b 100644
--- a/src/nbi/utils.py
+++ b/src/nbi/utils.py
@@ -34,6 +34,7 @@ def add_noise(x_err, x, y=None):
     x: light curve of shape (length,)
     y: parameter of shape (dim,)
     """
+    # TODO: generalize to multi-channel / multi-modal x (e.g., list of arrays)
     rand = np.random.normal(0, 1, size=x.shape[0])
     x_noise = x + rand * x_err
     return x_noise, y
@@ -68,54 +69,89 @@ def collate_list_channels(batch):
     return x_out, y_out
 
 
-def masked_mean_std(xs, mask_dim, masked_dims):
+def masked_mean_std(xs, mask_dim, masked_dims, reduce_axes=(0, 2), eps=1e-8):
     """
-    xs: np.ndarray [N, C, L]
-    mask_dim: int
-    masked_dims: list[int]
-    Returns mu, sd shaped [1, C, 1]
+    Compute masked mean/std for selected channels.
+
+    Parameters
+    ----------
+    xs : np.ndarray, shape [N, C, L]
+    mask_dim : int
+        Channel index containing the mask (same length L).
+    masked_dims : list[int]
+        Channels to compute masked stats for.
+    reduce_axes : tuple[int]
+        Axes to reduce over when computing stats.
+        - (0,2): reduce over batch and length -> outputs [1, C, 1]
+        - (0,):  reduce over batch only        -> outputs [1, C, L]
+    eps : float
+        Small value to avoid division by zero.
+
+    Returns
+    -------
+    mu, sd : np.ndarray
+        Shapes follow keepdims=True reduction:
+        - reduce_axes=(0,2) => [1, C, 1]
+        - reduce_axes=(0,)  => [1, C, L]
     """
-    C = xs.shape[1]
-    mu = np.zeros((1, C, 1), dtype=np.float32)
-    sd = np.ones((1, C, 1), dtype=np.float32)
+    xs = np.asarray(xs)
+    if xs.ndim != 3:
+        raise ValueError(f"xs must be [N,C,L], got shape {xs.shape}")
 
-    mask = xs[:, mask_dim, :]  # [N, L]
-    # force mask to {0,1} and finite
+    N, C, L = xs.shape
+    reduce_axes = tuple(reduce_axes)
+
+    # Determine expected output shape with keepdims=True
+    out_shape = [N, C, L]
+    for ax in reduce_axes:
+        out_shape[ax] = 1
+    out_shape = tuple(out_shape)
+
+    mu = np.zeros(out_shape, dtype=np.float32)
+    sd = np.ones(out_shape, dtype=np.float32)
+
+    # mask: [N, 1, L] broadcastable across channels
+    mask = xs[:, mask_dim:mask_dim + 1, :]  # [N,1,L]
     mask = np.where(np.isfinite(mask), mask, 0.0)
     mask = (mask > 0.5).astype(np.float32)
 
     for d in masked_dims:
-        v = xs[:, d, :]
-        # treat non-finite values as missing
-        finite = np.isfinite(v)
-        w = mask * finite.astype(np.float32)
-        wsum = w.sum()
-
-        if wsum <= 0:
-            # fallback: unmasked finite-only stats
-            vv = v[finite]
-            if vv.size == 0:
-                mu[0, d, 0] = 0.0
-                sd[0, d, 0] = 1.0
-                continue
-            m = float(vv.mean())
-            s = float(vv.std())
-            if (not np.isfinite(s)) or s == 0:
-                s = 1.0
-            mu[0, d, 0] = m
-            sd[0, d, 0] = s
-            continue
-
-        m = (v * w).sum() / wsum
-        var = ((v - m) ** 2 * w).sum() / wsum
-        s = np.sqrt(var)
+        v = xs[:, d:d + 1, :]  # [N,1,L]
+        finite = np.isfinite(v).astype(np.float32)
 
-        if not np.isfinite(m):
-            m = 0.0
-        if not np.isfinite(s) or s == 0:
-            s = 1.0
+        w = mask * finite  # [N,1,L], 0/1 weights
+        wsum = w.sum(axis=reduce_axes, keepdims=True)  # out_shape but with channel=1
+
+        # Safe denom
+        denom = np.maximum(wsum, 0.0)
+
+        # Masked mean
+        v_filled = np.where(np.isfinite(v), v, 0.0)
+        m = (v_filled * w).sum(axis=reduce_axes, keepdims=True) / np.maximum(denom, eps)
+
+        # Masked variance
+        var = ((v_filled - m) ** 2 * w).sum(axis=reduce_axes, keepdims=True) / np.maximum(denom, eps)
+        s = np.sqrt(var)
 
-        mu[0, d, 0] = float(m)
-        sd[0, d, 0] = float(s)
+        # Where denom==0, fall back to finite-only unmasked stats over the same reduce_axes
+        # This fallback is done elementwise in the reduced shape.
+        if np.any(denom <= 0):
+            w2 = finite
+            w2sum = w2.sum(axis=reduce_axes, keepdims=True)
+            m2 = (v_filled * w2).sum(axis=reduce_axes, keepdims=True) / np.maximum(w2sum, eps)
+            var2 = ((v_filled - m2) ** 2 * w2).sum(axis=reduce_axes, keepdims=True) / np.maximum(w2sum, eps)
+            s2 = np.sqrt(var2)
+
+            use_fallback = (denom <= 0)
+            m = np.where(use_fallback, m2, m)
+            s = np.where(use_fallback, s2, s)
+
+        # Final cleanup
+        m = np.where(np.isfinite(m), m, 0.0)
+        s = np.where(np.isfinite(s) & (s > 0), s, 1.0)
+
+        # write into channel d (broadcast along reduced axes)
+        mu[:, d:d + 1, :] = m.astype(np.float32)
+        sd[:, d:d + 1, :] = s.astype(np.float32)
 
     return mu, sd
\ No newline at end of file

From 1b96c409841ea4740891e178863c624bb1770951 Mon Sep 17 00:00:00 2001
From: Jackie <jackie.blaum@gmail.com>
Date: Wed, 11 Mar 2026 13:42:07 -0700
Subject: [PATCH 08/11] Add multi-modal test and fix bugs exposed by
 multi-channel path

- Add test_multi_modal exercising multi-channel Dataset, MultiFeaturizer,
  DatasetContainer, and save/load round-trip
- Fix missing torch.nn import in engine
- Fix fit() assertion to accept TorchDataset input
- Fix __init__ to accept dict state_dict (not just file paths)
- Fix log_prob() to use get_network() instead of self.network.module
- Fix collate_list_channels to handle both single-tensor and list inputs
---
 src/nbi/engine.py   |  19 +++++----
 src/nbi/utils.py    |  35 +++++++++++----
 test/test_engine.py | 102 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 139 insertions(+), 17 deletions(-)

diff --git a/src/nbi/engine.py b/src/nbi/engine.py
index 2d00803..58d141d 100644
--- a/src/nbi/engine.py
+++ b/src/nbi/engine.py
@@ -12,7 +12,7 @@
 import torch
 import wandb
 from multiprocess import Pool
-from torch import optim
+from torch import nn, optim
 from torch.optim.lr_scheduler import (
     CosineAnnealingWarmRestarts,
     MultiStepLR,
@@ -152,9 +152,12 @@ def __init__(
         self.init_env()
 
         if state_dict is not None:
-            _state_dict = torch.load(
-                state_dict, map_location=self.device, weights_only=False
-            )
+            if isinstance(state_dict, str):
+                _state_dict = torch.load(
+                    state_dict, map_location=self.device, weights_only=False
+                )
+            else:
+                _state_dict = state_dict
             flow_config_all = _state_dict["flow_config"]
             featurizer = (
                 featurizer
@@ -379,8 +382,8 @@ def fit(
         -------
 
         """
-        # either simulate n_sims or provide pre computed samples
-        assert n_sims > 0 or y is not None
+        # either simulate n_sims, provide pre computed samples, or pass a Dataset
+        assert n_sims > 0 or y is not None or isinstance(x, TorchDataset)
 
         if type(noise) == np.ndarray:
             # for i.i.d. gaussian noise
@@ -1638,8 +1641,8 @@ def log_prob(self, x, y):
         x = torch.from_numpy(x).to(self.device, dtype=torch.float32)
         y = torch.from_numpy(y).to(self.device, dtype=torch.float32)
         with torch.no_grad():
-            # it appears that DataParallal doesn't work properly here
-            log_prob = self.network.module(x, y).cpu().numpy()[:, 0] * -1
+            # use get_network() to handle both DataParallel and non-wrapped cases
+            log_prob = self.get_network()(x, y).cpu().numpy()[:, 0] * -1
         return log_prob
 
     def corner(
diff --git a/src/nbi/utils.py b/src/nbi/utils.py
index 93b228b..9986355 100644
--- a/src/nbi/utils.py
+++ b/src/nbi/utils.py
@@ -56,16 +56,33 @@ def log_like_iidg(x_err):
 
 def collate_list_channels(batch):
     """
-    Batch is a list of tuples: [(x_list, y), ...]
-    where x_list is list of tensors: [x0, x1, ..., xK]
-    Returns:
-    x_batched: list of tensors, each stacked to [B, ...]
-    y_batched: tensor [B, D]
+    Custom collate function that supports both single-tensor and multi-channel inputs.
+
+    For multi-channel data, batch is a list of tuples: [(x_list, y), ...]
+    where x_list is a list/tuple of tensors, one per channel: [x0, x1, ..., xK].
+
+    For single-tensor data, batch is a list of tuples: [(x, y), ...]
+    where x is a tensor or ndarray.
+
+    Returns
+    -------
+    x_batched : list of tensors (multi-channel) or single tensor (single-channel)
+        Each channel stacked to [B, ...], or a single stacked tensor [B, ...].
+    y_batched : tensor [B, D]
     """
-    xs, ys = zip(*batch)  # xs is tuple of lists
-    k = len(xs[0])
-    x_out = [torch.stack([x_i[j] for x_i in xs], dim=0) for j in range(k)]
-    y_out = torch.stack(list(ys), dim=0)
+    xs, ys = zip(*batch)
+
+    # Convert y to tensors
+    y_out = torch.stack([torch.as_tensor(y) for y in ys], dim=0)
+
+    # Multi-channel: x is a list/tuple of tensors
+    if isinstance(xs[0], (list, tuple)):
+        k = len(xs[0])
+        x_out = [torch.stack([torch.as_tensor(x_i[j]) for x_i in xs], dim=0) for j in range(k)]
+    else:
+        # Single tensor/ndarray
+        x_out = torch.stack([torch.as_tensor(x) for x in xs], dim=0)
+
     return x_out, y_out
 
 
diff --git a/test/test_engine.py b/test/test_engine.py
index e9bdcd2..ca9cd03 100644
--- a/test/test_engine.py
+++ b/test/test_engine.py
@@ -6,8 +6,10 @@
 import nbi
 import numpy as np
 import pytest
+import torch
 from scipy.stats import uniform
 from torch import nn
+from torch.utils.data import Dataset
 
 # Common setup variables
 t = np.linspace(0, 1, 50)
@@ -171,5 +173,105 @@ def test_custom_featurizer():
     fit_and_predict_anpe(engine)
 
 
+class MultiChannelSineDataset(Dataset):
+    """Synthetic dataset returning two channels per sample: (sine, cosine)."""
+
+    def __init__(self, n_samples=500, seq_len=50, seed=0):
+        rng = np.random.default_rng(seed)
+        self.t = np.linspace(0, 1, seq_len).astype(np.float32)
+        self.params = np.column_stack([
+            rng.uniform(0, 2 * np.pi, n_samples),   # phi0
+            rng.uniform(1, 5, n_samples),            # A
+            rng.uniform(2 * np.pi, 12 * np.pi, n_samples),  # omega
+        ]).astype(np.float32)
+
+        # Pre-compute both channels
+        self.ch0 = []  # sine channel
+        self.ch1 = []  # cosine channel
+        for phi0, A, omega in self.params:
+            self.ch0.append(A * np.sin(omega * self.t + phi0))
+            self.ch1.append(A * np.cos(omega * self.t + phi0))
+
+    def __len__(self):
+        return len(self.params)
+
+    def __getitem__(self, idx):
+        # Each channel: [1, seq_len]  (C=1 per channel)
+        x0 = torch.tensor(self.ch0[idx][None, :], dtype=torch.float32)
+        x1 = torch.tensor(self.ch1[idx][None, :], dtype=torch.float32)
+        y = torch.tensor(self.params[idx], dtype=torch.float32)
+        return [x0, x1], y
+
+
+def test_multi_modal():
+    """Test multi-channel input with per-channel featurizers and DatasetContainer."""
+    dim_out = 32
+    flow = {
+        "n_dims": 3,
+        "flow_hidden": 32,
+        "num_blocks": 4,
+        "num_cond_inputs": dim_out * 2,
+    }
+
+    featurizer_ch0 = nn.Sequential(
+        nn.Flatten(start_dim=1),
+        nn.Linear(50, 64),
+        nn.ReLU(),
+        nn.Linear(64, dim_out),
+    )
+    featurizer_ch0.num_outputs = dim_out
+
+    featurizer_ch1 = nn.Sequential(
+        nn.Flatten(start_dim=1),
+        nn.Linear(50, 64),
+        nn.ReLU(),
+        nn.Linear(64, dim_out),
+    )
+    featurizer_ch1.num_outputs = dim_out
+
+    dataset = MultiChannelSineDataset(n_samples=500, seq_len=50)
+
+    engine = nbi.NBI(
+        flow=flow,
+        featurizer=[featurizer_ch0, featurizer_ch1],
+        labels=labels,
+        path="test_multimodal",
+        device="cpu",
+    )
+
+    engine.fit(
+        x=dataset,
+        n_sims=-1,
+        n_rounds=1,
+        n_epochs=2,
+        batch_size=32,
+        lr=0.001,
+        min_lr=0.001,
+        workers=0,
+        plot=False,
+    )
+
+    # Build a single observation from both channels
+    y_test = np.array([1.0, 2.0, 6.0], dtype=np.float32)
+    x0_obs = (2.0 * np.sin(6.0 * t + 1.0)).astype(np.float32)[None, :]  # [1, L]
+    x1_obs = (2.0 * np.cos(6.0 * t + 1.0)).astype(np.float32)[None, :]
+    x_obs_multi = [x0_obs, x1_obs]
+
+    samples = engine.predict(x_obs_multi, n_samples=100, seed=0)
+    assert samples.shape == (100, 3)
+
+    # Test save/load round-trip
+    best_params = engine.get_params()
+    engine2 = nbi.NBI(
+        state_dict=best_params,
+        featurizer=[featurizer_ch0, featurizer_ch1],
+        labels=labels,
+        path="test_multimodal2",
+        device="cpu",
+    )
+    samples2 = engine2.predict(x_obs_multi, n_samples=100, seed=0)
+    assert np.allclose(samples, samples2)
+
+
 if __name__ == "__main__":
     pytest.main()

From 443a242b2a07868fcff46a8d711732f4f3f57ec2 Mon Sep 17 00:00:00 2001
From: Jackie <jackie.blaum@gmail.com>
Date: Wed, 11 Mar 2026 13:44:39 -0700
Subject: [PATCH 09/11] Add test artifacts to .gitignore

---
 .gitignore | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/.gitignore b/.gitignore
index bb16b6d..ffb12e6 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,5 +7,10 @@ src/nbi.egg-info/
 
 *.DS_Store
 
+# Test artifacts
+test/*.npy
+test/[0-9]*/
+test_multimodal/
+
 # Sphinx documentation
 docs/_build/

From f5ee6ddf65c15fcf35a42c8e6459317528517ab4 Mon Sep 17 00:00:00 2001
From: Jackie <jackie.blaum@gmail.com>
Date: Mon, 23 Mar 2026 23:52:08 -0700
Subject: [PATCH 10/11] Fix EmpiricalPrior.logpdf() for list-of-scipy priors,
 restore importance_reweight ValueError

EmpiricalPrior.logpdf() assumed self.priors was a single object with a
.logpdf() method, but in practice it is typically a list of per-parameter
scipy distributions. This broke SNPE (which calls log_prior via
importance_reweight) while ANPE was unaffected since it never reaches
that code path during training.

Also restores ValueError in importance_reweight when all weights are
invalid (was silently returning zeros), and fixes the corresponding test
which was never reaching the raise because engine.like was None.
---
 src/nbi/empirical_prior.py |  10 ++-
 src/nbi/engine.py          |  38 ++++++--
 test/test_engine.py        | 179 +++++++++++++++++++++++++++++++++++++
 3 files changed, 220 insertions(+), 7 deletions(-)

diff --git a/src/nbi/empirical_prior.py b/src/nbi/empirical_prior.py
index ec22946..5b480af 100644
--- a/src/nbi/empirical_prior.py
+++ b/src/nbi/empirical_prior.py
@@ -19,6 +19,14 @@ def rvs(self, size=1, replace=True):
     def logpdf(self, params):
         if self.priors is None:
             raise NotImplementedError(
-                "logp not implemented: provide a prior model or override this method."
+                "logpdf not implemented: provide a `priors` argument (a list of "
+                "scipy-style distributions or a single object with a .logpdf method) "
+                "to enable log-probability evaluation, which is required for SNPE."
             )
+        params = np.asarray(params)
+        if isinstance(self.priors, list):
+            log_prob = np.zeros(len(params))
+            for i, prior in enumerate(self.priors):
+                log_prob += prior.logpdf(params[:, i])
+            return log_prob
         return self.priors.logpdf(params)
diff --git a/src/nbi/engine.py b/src/nbi/engine.py
index 58d141d..16b9103 100644
--- a/src/nbi/engine.py
+++ b/src/nbi/engine.py
@@ -32,6 +32,14 @@
 from .empirical_prior import EmpiricalPrior
 from .utils import iid_gaussian, log_like_iidg, parallel_simulate, collate_list_channels, masked_mean_std
 
+
+def _worker_init_fn(worker_id):
+    """Prevent numpy/BLAS thread oversubscription in DataLoader workers."""
+    os.environ["OMP_NUM_THREADS"] = "1"
+    os.environ["MKL_NUM_THREADS"] = "1"
+    os.environ["OPENBLAS_NUM_THREADS"] = "1"
+
+
 class NBI:
     """Neural Bayesian Inference Engine.
 
@@ -689,8 +697,11 @@ def importance_reweight(self, x_obs, x, y):
         bad = np.isnan(log_weights) + np.isinf(log_weights)
         valid_weights = log_weights[~bad]
         if len(valid_weights) == 0:
-            print("All log weights are NaN or Inf — skipping this round!")
-            return np.zeros_like(log_weights)
+            raise ValueError(
+                "All importance weights are invalid (NaN or Inf). "
+                "This typically means the prior, likelihood, or proposal "
+                "returned degenerate values for every sample."
+            )
         log_weights -= log_weights[~bad].max()
 
         weights = np.exp(log_weights)
@@ -1430,15 +1441,30 @@ def _init_loader(self, data_container, batch_size, workers=4, pin_memory=None, p
             "drop_last": True,
             "persistent_workers": True,
             "prefetch_factor": prefetch_factor if workers > 0 else None,
-            "collate_fn": collate_list_channels, 
+            "collate_fn": collate_list_channels,
+            "worker_init_fn": _worker_init_fn,
         }
         if workers == 0:
             kwargs.pop("prefetch_factor", None)
             kwargs["persistent_workers"] = False
 
-        self.train_loader = DataLoader(
-            train_container, batch_size=batch_size, shuffle=True, **kwargs
-        )
+        # Use ShardGroupedSampler when the underlying dataset supports it,
+        # to minimise shard-switch overhead with mmap-backed .npy shards.
+        train_sampler = None
+        try:
+            from ebsbi.shards import ShardGroupedSampler
+            train_sampler = ShardGroupedSampler(train_container)
+        except (ImportError, TypeError):
+            pass
+
+        if train_sampler is not None:
+            self.train_loader = DataLoader(
+                train_container, batch_size=batch_size, sampler=train_sampler, **kwargs
+            )
+        else:
+            self.train_loader = DataLoader(
+                train_container, batch_size=batch_size, shuffle=True, **kwargs
+            )
         self.valid_loader = DataLoader(val_container, batch_size=batch_size, **kwargs)
 
         # if self.network_reinit or self.round == 0:
diff --git a/test/test_engine.py b/test/test_engine.py
index ca9cd03..39a3d90 100644
--- a/test/test_engine.py
+++ b/test/test_engine.py
@@ -10,6 +10,7 @@
 from scipy.stats import uniform
 from torch import nn
 from torch.utils.data import Dataset
+from nbi.empirical_prior import EmpiricalPrior
 
 # Common setup variables
 t = np.linspace(0, 1, 50)
@@ -273,5 +274,183 @@ def test_multi_modal():
     assert np.allclose(samples, samples2)
 
 
+def test_importance_reweight_raises_when_all_log_weights_invalid():
+    flow = {
+        "n_dims": 3,
+        "flow_hidden": 16,
+        "num_blocks": 2,
+    }
+
+    featurizer = {
+        "type": "resnet-gru",
+        "norm": "weight_norm",
+        "dim_in": 1,
+        "dim_out": 16,
+        "dim_conv_max": 32,
+        "depth": 2,
+    }
+
+    engine = nbi.NBI(
+        flow=flow,
+        featurizer=featurizer,
+        simulator=sine,
+        priors=priors,
+        labels=labels,
+        path="test_invalid_weights",
+        device="cpu",
+        n_jobs=1,
+    )
+
+    engine.like = True  # must be non-None to enter the reweight logic
+    engine.log_like = lambda x_obs, x, y: np.full(len(y), np.nan)
+    engine.log_prob = lambda x_obs, y: np.zeros(len(y))
+    y = np.zeros((5, 3), dtype=np.float32)
+
+    with pytest.raises(ValueError, match="All importance weights are invalid"):
+        engine.importance_reweight(np.zeros(50), np.zeros((5, 50)), y)
+
+
+def test_empirical_prior_logpdf_with_list_priors():
+    """EmpiricalPrior.logpdf must handle a list of per-parameter scipy distributions."""
+    lookup = np.column_stack([
+        uniform(loc=0, scale=2 * np.pi).rvs(100),
+        uniform(loc=1, scale=4).rvs(100),
+        uniform(loc=2 * np.pi, scale=10 * np.pi).rvs(100),
+    ])
+    ep = EmpiricalPrior(lookup, priors=priors)
+
+    # Samples inside the prior support should give finite log-probabilities
+    params_in = ep.rvs(10)
+    lp = ep.logpdf(params_in)
+    assert lp.shape == (10,)
+    assert np.all(np.isfinite(lp))
+
+    # A point outside every prior's support should give -inf
+    params_out = np.array([[-999, -999, -999]])
+    lp_out = ep.logpdf(params_out)
+    assert np.all(np.isinf(lp_out))
+
+
+def test_empirical_prior_logpdf_raises_without_priors():
+    """EmpiricalPrior without priors must raise NotImplementedError on logpdf."""
+    lookup = np.random.randn(50, 3)
+    ep = EmpiricalPrior(lookup)
+    with pytest.raises(NotImplementedError, match="logpdf not implemented"):
+        ep.logpdf(np.random.randn(5, 3))
+
+
+def test_snpe_with_empirical_prior():
+    """SNPE (n_rounds > 1) must work when the prior is an EmpiricalPrior with list priors."""
+    # Build lookup table by sampling from the analytic priors
+    np.random.seed(42)
+    lookup = np.column_stack([p.rvs(500) for p in priors])
+    emp_prior = EmpiricalPrior(lookup, priors=priors, random_state=42)
+
+    flow = {
+        "n_dims": 3,
+        "flow_hidden": 32,
+        "num_blocks": 4,
+    }
+
+    featurizer = {
+        "type": "resnet-gru",
+        "norm": "weight_norm",
+        "dim_in": 1,
+        "dim_out": 32,
+        "dim_conv_max": 256,
+        "depth": 3,
+    }
+
+    engine = nbi.NBI(
+        flow=flow,
+        featurizer=featurizer,
+        simulator=sine,
+        priors=emp_prior,
+        labels=labels,
+        path="test_empirical_snpe",
+        device="cpu",
+        n_jobs=4,
+    )
+
+    # Run 2 rounds of SNPE — this exercises log_prior via importance_reweight
+    engine.fit(
+        x_obs=x_obs,
+        y_true=y_true,
+        n_sims=320,
+        n_rounds=2,
+        n_epochs=100,
+        batch_size=32,
+        lr=0.001,
+        min_lr=0.001,
+        early_stop_train=True,
+        early_stop_patience=1,
+        noise=np.array([1] * 50),
+        workers=10,
+        plot=False,
+    )
+
+    y, w = engine.predict(
+        x_obs,
+        x_err=np.array([0.2] * 50),
+        y_true=y_true,
+        n_samples=1000,
+        neff_min=100,
+        f_accept_min=0.1,
+        seed=0,
+    )
+
+    assert y.shape[1] == 3
+    assert w is not None
+
+
+def test_anpe_with_empirical_prior():
+    """ANPE (n_rounds=1) must work with EmpiricalPrior with list priors."""
+    np.random.seed(42)
+    lookup = np.column_stack([p.rvs(500) for p in priors])
+    # priors needed because predict() calls _draw_params() → log_prior()
+    emp_prior = EmpiricalPrior(lookup, priors=priors, random_state=42)
+
+    flow = {
+        "n_dims": 3,
+        "flow_hidden": 32,
+        "num_blocks": 4,
+    }
+
+    featurizer = {
+        "type": "resnet-gru",
+        "norm": "weight_norm",
+        "dim_in": 1,
+        "dim_out": 32,
+        "dim_conv_max": 256,
+        "depth": 3,
+    }
+
+    engine = nbi.NBI(
+        flow=flow,
+        featurizer=featurizer,
+        simulator=sine,
+        priors=emp_prior,
+        labels=labels,
+        path="test_empirical_anpe",
+        device="cpu",
+        n_jobs=4,
+    )
+
+    engine.fit(
+        x_obs=x_obs,
+        n_sims=320,
+        n_rounds=1,
+        n_epochs=1,
+        batch_size=32,
+        lr=0.001,
+        min_lr=0.001,
+        workers=10,
+        plot=False,
+    )
+
+    y = engine.predict(x_obs, n_samples=1000, seed=0)
+    assert len(y) == 1000
+
+
 if __name__ == "__main__":
     pytest.main()

From 92684870a0e98f1168cfdb22d4bdfa798780257f Mon Sep 17 00:00:00 2001
From: Jackie <jackie.blaum@gmail.com>
Date: Thu, 9 Apr 2026 22:20:42 -0700
Subject: [PATCH 11/11] Adopt upstream behavior: return zeros instead of
 raising on all-invalid weights

Aligns importance_reweight with upstream PR #40, which silently skips a
round when every weight is NaN/Inf rather than crashing. Updates the
test accordingly to assert zero weights are returned.
---
 src/nbi/engine.py   | 7 ++-----
 test/test_engine.py | 8 +++++---
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/src/nbi/engine.py b/src/nbi/engine.py
index 16b9103..406eb3e 100644
--- a/src/nbi/engine.py
+++ b/src/nbi/engine.py
@@ -697,11 +697,8 @@ def importance_reweight(self, x_obs, x, y):
         bad = np.isnan(log_weights) + np.isinf(log_weights)
         valid_weights = log_weights[~bad]
         if len(valid_weights) == 0:
-            raise ValueError(
-                "All importance weights are invalid (NaN or Inf). "
-                "This typically means the prior, likelihood, or proposal "
-                "returned degenerate values for every sample."
-            )
+            print("All log weights are NaN or Inf — skipping this round!")
+            return np.zeros_like(log_weights)
         log_weights -= log_weights[~bad].max()
 
         weights = np.exp(log_weights)
diff --git a/test/test_engine.py b/test/test_engine.py
index 39a3d90..2134d70 100644
--- a/test/test_engine.py
+++ b/test/test_engine.py
@@ -274,7 +274,7 @@ def test_multi_modal():
     assert np.allclose(samples, samples2)
 
 
-def test_importance_reweight_raises_when_all_log_weights_invalid():
+def test_importance_reweight_returns_zeros_when_all_log_weights_invalid():
     flow = {
         "n_dims": 3,
         "flow_hidden": 16,
@@ -306,8 +306,10 @@ def test_importance_reweight_raises_when_all_log_weights_invalid():
     engine.log_prob = lambda x_obs, y: np.zeros(len(y))
     y = np.zeros((5, 3), dtype=np.float32)
 
-    with pytest.raises(ValueError, match="All importance weights are invalid"):
-        engine.importance_reweight(np.zeros(50), np.zeros((5, 50)), y)
+    weights = engine.importance_reweight(np.zeros(50), np.zeros((5, 50)), y)
+    assert weights is not None
+    assert weights.shape == (5,)
+    assert np.all(weights == 0)
 
 
 def test_empirical_prior_logpdf_with_list_priors():