cweniger
diff --git a/‎examples/05_linear_regression/config.yaml‎
Lines changed: 2 additions & 0 deletions b/‎examples/05_linear_regression/config.yaml‎
Lines changed: 2 additions & 0 deletions
diff --git a/‎falcon/cli.py‎
Lines changed: 23 additions & 15 deletions b/‎falcon/cli.py‎
Lines changed: 23 additions & 15 deletions
diff --git a/‎falcon/contrib/SNPE_A.py‎
Lines changed: 12 additions & 13 deletions b/‎falcon/contrib/SNPE_A.py‎
Lines changed: 12 additions & 13 deletions
diff --git a/‎falcon/contrib/stepwise_estimator.py‎
Lines changed: 21 additions & 21 deletions b/‎falcon/contrib/stepwise_estimator.py‎
Lines changed: 21 additions & 21 deletions
diff --git a/‎falcon/core/base_estimator.py‎
Lines changed: 7 additions & 8 deletions b/‎falcon/core/base_estimator.py‎
Lines changed: 7 additions & 8 deletions
@@ -110,6 +110,7 @@ graph:
 
     ray:
       num_gpus: 0.5
+      chunk_size: 128
 
   y:
     parents: [theta]
@@ -120,6 +121,7 @@ graph:
     observed: "./data/mock_data.npz['y']"
     ray:
       num_gpus: 0.5
+      chunk_size: 128
 
 # -----------------------------------------------------------------------------
 # Sampling configuration
 
@@ -411,46 +411,54 @@ def sample_mode(cfg, sample_type: str) -> None:
 
     if sample_type == "prior":
         # Generate forward samples from prior
-        samples = deployed_graph.sample(num_samples)
+        sample_refs = deployed_graph.sample(num_samples)
 
     elif sample_type == "posterior":
-        # TODO: Implement posterior sampling (requires trained model and observations)
         deployed_graph.load(Path(cfg.paths.graph))
-        samples = deployed_graph.sample_posterior(num_samples, observations)
+        sample_refs = deployed_graph.sample_posterior(num_samples, observations)
 
     elif sample_type == "proposal":
-        # Proposal sampling requires observations for conditioning
-        # Load observations from config
         deployed_graph.load(Path(cfg.paths.graph))
-        samples = deployed_graph.sample_proposal(num_samples, observations)
+        sample_refs = deployed_graph.sample_proposal(num_samples, observations)
 
     else:
         raise ValueError(f"Unknown sample type: {sample_type}")
 
-    # Apply smart key selection based on mode and user overrides
+    # Resolve refs to arrays (keys are flat: 'theta.value', 'theta.log_prob', 'x.value', ...)
+    samples = deployed_graph._refs_to_arrays(sample_refs)
+
+    # Build key selection based on node names (strip .value/.log_prob suffixes)
+    node_keys = {k for k in samples.keys() if k.endswith('.value')}
+
     if sample_type in ["prior", "proposal"]:
-        # Default: save everything
-        default_keys = set(samples.keys())
+        # Default: save all .value keys
+        default_keys = set(node_keys)
     elif sample_type == "posterior":
         # Default: save only posterior nodes (nodes with evidence)
         default_keys = {
-            k for k, node in graph.node_dict.items() if node.evidence and k in samples
+            f"{k}.value" for k, node in graph.node_dict.items()
+            if node.evidence and f"{k}.value" in samples
         }
 
-    # Apply user overrides
+    # Apply user overrides (user specifies node names, we match .value keys)
     exclude_keys = sample_cfg.get("exclude_keys", None)
     add_keys = sample_cfg.get("add_keys", None)
 
     if exclude_keys:
-        exclude_set = set(exclude_keys.split(","))
+        exclude_set = {f"{k}.value" for k in exclude_keys.split(",")}
         default_keys -= exclude_set
 
     if add_keys:
-        add_set = set(add_keys.split(","))
+        add_set = {f"{k}.value" for k in add_keys.split(",")}
         default_keys |= add_set
 
-    # Filter samples to selected keys
-    save_data = {k: samples[k] for k in default_keys if k in samples}
+    # Filter samples to selected keys, strip .value suffix for user-facing output
+    save_data = {}
+    for k in default_keys:
+        if k in samples:
+            # Strip '.value' suffix for cleaner output key names
+            user_key = k[:-6] if k.endswith('.value') else k
+            save_data[user_key] = samples[k]
 
     print(f"Generated samples with shapes:")
     for key, value in save_data.items():
 
@@ -12,7 +12,6 @@
 from torch.optim import AdamW
 from torch.optim.lr_scheduler import ReduceLROnPlateau
 
-from falcon.core.utils import RVBatch
 from falcon.core.logger import log, debug, info, warning, error
 from falcon.contrib.flow import Flow
 from falcon.contrib.stepwise_estimator import StepwiseEstimator, TrainingLoopConfig
@@ -228,10 +227,10 @@ def _unpack_batch(self, batch, phase: str):
             Tuple of (ids, theta, theta_logprob, conditions, u, u_device, conditions_device)
         """
         ids = batch._ids
-        theta = self._to_tensor(batch[self.theta_key])
-        theta_logprob = self._to_tensor(batch[f"{self.theta_key}.logprob"])
+        theta = self._to_tensor(batch[f"{self.theta_key}.value"])
+        theta_logprob = self._to_tensor(batch[f"{self.theta_key}.log_prob"])
         conditions = {
-            k: self._to_tensor(batch[k]) for k in self.condition_keys if k in batch
+            k: self._to_tensor(batch[f"{k}.value"]) for k in self.condition_keys if f"{k}.value" in batch
         }
 
         # Record IDs for history
@@ -345,34 +344,34 @@ def on_epoch_end(self, epoch: int, val_metrics: Dict[str, float]) -> Optional[Di
 
     # ==================== Sampling Methods ====================
 
-    def sample_prior(self, num_samples: int, conditions: Optional[Dict] = None) -> RVBatch:
+    def sample_prior(self, num_samples: int, conditions: Optional[Dict] = None) -> dict:
         """Sample from the prior distribution."""
         if conditions:
             raise ValueError("Conditions are not supported for sample_prior.")
         samples = self.simulator_instance.simulate_batch(num_samples)
         # Log probability for uniform prior over hypercube [-bound, bound]^d
         bound = self.config.inference.hypercube_bound
-        logprob = np.ones(num_samples) * (-np.log(2 * bound) ** self.param_dim)
-        return RVBatch(samples, logprob=logprob)
+        log_prob = np.ones(num_samples) * (-np.log(2 * bound) ** self.param_dim)
+        return {'value': samples, 'log_prob': log_prob}
 
     def sample_posterior(
         self,
         num_samples: int,
         conditions: Optional[Dict] = None,
-    ) -> RVBatch:
+    ) -> dict:
         """Sample from the posterior distribution q(theta|x)."""
         # Fall back to prior if networks not yet initialized (training hasn't started)
         if not self.networks_initialized:
             return self.sample_prior(num_samples)
 
         samples, logprob = self._importance_sample(num_samples, mode="posterior", conditions=conditions or {})
-        return RVBatch(samples.numpy(), logprob=logprob.numpy())
+        return {'value': samples.numpy(), 'log_prob': logprob.numpy()}
 
     def sample_proposal(
         self,
         num_samples: int,
         conditions: Optional[Dict] = None,
-    ) -> RVBatch:
+    ) -> dict:
         """Sample from the widened proposal distribution for adaptive resampling."""
         # Fall back to prior if networks not yet initialized (training hasn't started)
         if not self.networks_initialized:
@@ -393,7 +392,7 @@ def sample_proposal(
             "sample_proposal:std": samples.std().item(),
             "sample_proposal:logprob": logprob.mean().item(),
         })
-        return RVBatch(samples.numpy(), logprob=logprob.numpy())
+        return {'value': samples.numpy(), 'log_prob': logprob.numpy()}
 
     def _importance_sample(
         self,
@@ -405,8 +404,8 @@ def _importance_sample(
         cfg_inf = self.config.inference
 
         assert conditions, "Conditions must be provided."
-        # Move conditions to device
-        conditions = {k: v.to(self.device) for k, v in conditions.items()}
+        # Move conditions to device (handles both numpy arrays and tensors)
+        conditions = {k: self._to_tensor(v, self.device) for k, v in conditions.items()}
 
         # Use best models if available and configured, otherwise fall back to current
         use_best = cfg_inf.use_best_models_during_inference and self._best_conditional_flow is not None
 
@@ -15,7 +15,6 @@
 from torch.optim.lr_scheduler import ReduceLROnPlateau
 
 from falcon.core.base_estimator import BaseEstimator
-from falcon.core.utils import RVBatch
 from falcon.core.logger import log, debug, info, warning, error
 
 
@@ -117,7 +116,7 @@ def _to_tensor(x, device=None):
         """Convert numpy array or torch tensor to the target device."""
         if isinstance(x, torch.Tensor):
             return x if device is None else x.to(device)
-        return torch.from_numpy(x) if device is None else torch.from_numpy(x).to(device)
+        return torch.from_numpy(np.asarray(x)) if device is None else torch.from_numpy(np.asarray(x)).to(device)
 
     # ==================== Abstract Methods ====================
 
@@ -185,7 +184,8 @@ async def train(self, buffer) -> None:
             buffer: BufferView providing access to training/validation data
         """
         cfg = self.loop_config
-        keys = [self.theta_key, f"{self.theta_key}.logprob", *self.condition_keys]
+        keys = [f"{self.theta_key}.value", f"{self.theta_key}.log_prob",
+                *[f"{k}.value" for k in self.condition_keys]]
         await self._train(buffer, cfg, keys)
 
     async def _train(self, buffer, cfg, keys) -> None:
@@ -430,9 +430,9 @@ def _build_model(self, batch) -> nn.Module:
         from falcon.contrib.torch_embedding import instantiate_embedding
 
         # Extract and store tensors for reload
-        self._init_theta = self._to_tensor(batch[self.theta_key])
+        self._init_theta = self._to_tensor(batch[f"{self.theta_key}.value"])
         self._init_conditions = {
-            k: self._to_tensor(batch[k]) for k in self.condition_keys if k in batch
+            k: self._to_tensor(batch[f"{k}.value"]) for k in self.condition_keys if f"{k}.value" in batch
         }
         return self._create_model(self._init_theta, self._init_conditions)
 
@@ -471,11 +471,11 @@ def _create_model(self, theta: torch.Tensor, conditions: Dict[str, torch.Tensor]
 
     def _compute_loss(self, batch) -> Tuple[torch.Tensor, Dict[str, float]]:
         """Compute loss from batch."""
-        theta = self._to_tensor(batch[self.theta_key], self.device)
-        theta_logprob = self._to_tensor(batch[f"{self.theta_key}.logprob"])
+        theta = self._to_tensor(batch[f"{self.theta_key}.value"], self.device)
+        theta_logprob = self._to_tensor(batch[f"{self.theta_key}.log_prob"])
         conditions = {
-            k: self._to_tensor(batch[k], self.device)
-            for k in self.condition_keys if k in batch
+            k: self._to_tensor(batch[f"{k}.value"], self.device)
+            for k in self.condition_keys if f"{k}.value" in batch
         }
 
         # Transform theta to latent space if mode specified
@@ -601,49 +601,49 @@ def has_trained_model(self) -> bool:
 
     # ==================== Sampling Methods ====================
 
-    def sample_prior(self, num_samples: int, conditions: Optional[Dict] = None) -> RVBatch:
+    def sample_prior(self, num_samples: int, conditions: Optional[Dict] = None) -> dict:
         """Sample from the prior distribution."""
         if conditions:
             raise ValueError("Conditions are not supported for sample_prior.")
         samples = self.simulator_instance.simulate_batch(num_samples)
-        logprob = np.zeros(num_samples)
-        return RVBatch(samples, logprob=logprob)
+        log_prob = np.zeros(num_samples)
+        return {'value': samples, 'log_prob': log_prob}
 
-    def _sample(self, num_samples: int, conditions: Optional[Dict], gamma: Optional[float]) -> RVBatch:
+    def _sample(self, num_samples: int, conditions: Optional[Dict], gamma: Optional[float]) -> dict:
         """Internal sampling using inference model. Falls back to prior if not trained."""
         if not self.has_trained_model:
             return self.sample_prior(num_samples)
 
         assert conditions, "Conditions must be provided for sampling."
 
         conditions_device = {
-            k: v.to(self.device).expand(num_samples, *v.shape[1:])
+            k: self._to_tensor(v, self.device).expand(num_samples, *v.shape[1:])
             for k, v in conditions.items()
         }
 
         model = self.inference_model
         with torch.no_grad():
             model.eval()
             samples = model.sample(conditions_device, gamma=gamma)
-            logprob = model.log_prob(samples, conditions_device)
+            log_prob = model.log_prob(samples, conditions_device)
 
             # Transform samples from latent space back to theta space
             if self.latent_mode is not None:
                 samples = self.simulator_instance.forward(samples, mode=self.latent_mode)
 
-        return RVBatch(samples.cpu().numpy(), logprob=logprob.cpu().numpy())
+        return {'value': samples.cpu().numpy(), 'log_prob': log_prob.cpu().numpy()}
 
-    def sample_posterior(self, num_samples: int, conditions: Optional[Dict] = None) -> RVBatch:
+    def sample_posterior(self, num_samples: int, conditions: Optional[Dict] = None) -> dict:
         """Sample from the posterior distribution q(theta|x)."""
         return self._sample(num_samples, conditions, gamma=None)
 
-    def sample_proposal(self, num_samples: int, conditions: Optional[Dict] = None) -> RVBatch:
+    def sample_proposal(self, num_samples: int, conditions: Optional[Dict] = None) -> dict:
         """Sample from widened proposal distribution for adaptive resampling."""
         result = self._sample(num_samples, conditions, gamma=self.inference_config.gamma)
         log({
-            "sample_proposal:mean": result.value.mean(),
-            "sample_proposal:std": result.value.std(),
-            "sample_proposal:logprob": result.logprob.mean(),
+            "sample_proposal:mean": result['value'].mean(),
+            "sample_proposal:std": result['value'].std(),
+            "sample_proposal:logprob": result['log_prob'].mean(),
         })
         return result
 
 
@@ -6,8 +6,6 @@
 
 import torch
 
-from falcon.core.utils import RVBatch
-
 # Type alias for conditions: maps node names to tensors
 Conditions = Dict[str, torch.Tensor]
 
@@ -20,6 +18,7 @@ class BaseEstimator(ABC):
     Concrete implementations must provide all functionality.
 
     Conditions are passed as Dict[str, Tensor] mapping node names to values.
+    Sampling methods return dicts with 'value' (ndarray) and optionally 'log_prob' (ndarray).
     """
 
     @abstractmethod
@@ -35,7 +34,7 @@ async def train(self, buffer) -> None:
     @abstractmethod
     def sample_prior(
         self, num_samples: int, conditions: Optional[Conditions] = None
-    ) -> RVBatch:
+    ) -> dict:
         """
         Sample from the prior distribution.
 
@@ -44,14 +43,14 @@ def sample_prior(
             conditions: Conditioning values from parent nodes (usually None for prior)
 
         Returns:
-            RVBatch with samples and log probabilities
+            Dict with 'value' (ndarray) and optionally 'log_prob' (ndarray)
         """
         pass
 
     @abstractmethod
     def sample_posterior(
         self, num_samples: int, conditions: Optional[Conditions] = None
-    ) -> RVBatch:
+    ) -> dict:
         """
         Sample from the posterior distribution.
 
@@ -60,14 +59,14 @@ def sample_posterior(
             conditions: Dict mapping node names to condition tensors
 
         Returns:
-            RVBatch with samples and log probabilities
+            Dict with 'value' (ndarray) and optionally 'log_prob' (ndarray)
         """
         pass
 
     @abstractmethod
     def sample_proposal(
         self, num_samples: int, conditions: Optional[Conditions] = None
-    ) -> RVBatch:
+    ) -> dict:
         """
         Sample from the proposal distribution for adaptive resampling.
 
@@ -76,7 +75,7 @@ def sample_proposal(
             conditions: Dict mapping node names to condition tensors
 
         Returns:
-            RVBatch with samples and log probabilities
+            Dict with 'value' (ndarray) and optionally 'log_prob' (ndarray)
         """
         pass