Add self-contained implementation of the Simformer approach - #36
Add self-contained implementation of the Simformer approach#36tHarvey303 wants to merge 2 commits into
Conversation
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…iants Adds src/synference/simformer/, a self-contained PyTorch port of the original JAX Simformer (Gloeckler et al. 2024, scoresbibm/probjax), and rewires Simformer_Fitter to use it. Replaces (a) the scoresbibm/JAX wrapper (src/synference/simformer.py, deleted) and (b) SBI_Fitter.run_single_simformer, which depended on the unmerged sbi PR #1621 (deleted). jax is removed from the dependencies. The new subpackage faithfully replicates the paper architecture and math: per-variable tokens with frozen node-id embeddings and a mask-gated condition token, pre-LN transformer with time context injected in every MLP block, VE-SDE (default) and VP-SDE, masked denoising score matching with structured-random condition-mask sampling, adaptive gradient clipping, and early stopping. Inference supports arbitrary condition masks (posterior, likelihood, joint, missing bands), base attention masks, interval-constrained sampling via Tweedie guidance, and log probabilities via the probability-flow ODE. Models serialize to a plain torch payload and reload in a fresh process (SimformerModel.save/load). API notes (clean break): - sample_posterior(condition_mask=...) replaces the misnamed attention_mask argument (it was always a condition mask); full sentinel unchanged. - log_prob(X_test, theta, ...) takes theta as the second positional argument, fixing evaluate_model silently passing y_test as a mask. - Config overrides are validated against known keys instead of being silently swallowed. - Old JAX-era model pickles cannot be loaded by the new implementation. - fit_catalogue's Simformer branch now transposes samples like the regular branch, fixing wrong quantiles for multi-row catalogues. Validation: analytic 2D-Gaussian conditional gate, TwoMoons (bimodal posterior, TARP calibration, likelihood/joint conditionals, interval guidance, log-prob), and the galaxy test library end-to-end (training, artifacts, reload in a fresh process, fit_catalogue) in tests/test_simformer.py; full suite passes. A demo lives in examples/simformer/scripts/two_moons_demo.py and the docs notebook was rewritten for the new API. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
📝 WalkthroughWalkthroughThe PR replaces the legacy JAX Simformer integration with a native PyTorch implementation, including transformer and diffusion components, training, conditional sampling, interval guidance, log probabilities, serialization, SBI integration, examples, documentation, dependency updates, and comprehensive tests. ChangesNative Simformer foundations
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant TrainingData
participant train_simformer
participant SimformerModel
participant Simformer_Fitter
participant Sampling
TrainingData->>train_simformer: theta and x training arrays
train_simformer->>SimformerModel: trained score network and SDE
Simformer_Fitter->>SimformerModel: save or load native model
Simformer_Fitter->>Sampling: condition mask and observations
Sampling->>Simformer_Fitter: posterior, interval, or log-probability results
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (4)
src/synference/sbi_runner.py (1)
8397-8402: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer an exception over
assertfor this input contract.
sample_posterioris a public entry point and the column/mask check is user-facing validation;assertis removed underpython -O, after which a mismatch surfaces deeper insideSimformerModel._prepare_x_owith a less obvious message.♻️ Proposed change
- assert X_test.shape[1] == mask.sum(), ( - f"X_test has {X_test.shape[1]} columns but the condition mask expects " - f"{int(mask.sum())} observed values." - ) + if X_test.shape[1] != int(mask.sum()): + raise ValueError( + f"X_test has {X_test.shape[1]} columns but the condition mask expects " + f"{int(mask.sum())} observed values." + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/synference/sbi_runner.py` around lines 8397 - 8402, In sample_posterior, replace the assert guarding the X_test column count against mask.sum() with explicit user-facing validation that raises an appropriate exception containing the existing mismatch details. Preserve the current shape normalization and allow valid column/mask matches to proceed unchanged.tests/test_simformer.py (1)
405-427: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider marking the training fixtures as slow.
Both module-scoped fixtures train real models (3000 steps on 10k samples, plus 600 steps on the galaxy library) and downstream tests draw thousands of samples with 100 integration steps. That makes
pyteston the default selection substantially slower. A@pytest.mark.slowon these classes (with a CI opt-in) would keep the fast suite fast without losing coverage.Also applies to: 533-561
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_simformer.py` around lines 405 - 427, Mark the model-training test fixtures or their containing test classes, including the two-moons fixture two_moons_model and the galaxy-library fixture covering the other referenced range, with pytest.mark.slow so they are excluded from the default fast suite while remaining available through CI’s slow-test opt-in.src/synference/simformer/sampling.py (1)
179-203: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueExact divergence cost scales with the number of latent nodes.
One backward pass per latent dimension per ODE step (
num_steps × n_latentgraph traversals) makesdivergence="exact"expensive for galaxy models with many parameters, and it is the default inSimformerModel.log_probandSimformer_Fitter.log_prob. Consider defaulting to Hutchinson above a dimension threshold, or documenting the cost at the public entry points.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/synference/simformer/sampling.py` around lines 179 - 203, Address the scalability of exact divergence in _divergence_exact by either switching to Hutchinson divergence above a clearly defined latent-dimension threshold or documenting the num_steps × n_latent backward-pass cost at the public SimformerModel.log_prob and Simformer_Fitter.log_prob entry points. Preserve exact divergence behavior below any threshold and keep the existing log-probability API consistent.src/synference/simformer/model.py (1)
538-578: 🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick winHarden
SimformerModel.loadto avoid unpickling arbitrary objects.
weights_only=Falsekeeps full pickle execution available even though the saved payload is just model weights/config and numpy arrays. Useweights_only=Truewith explicit safe globals for the numpy dtypes/arrays needed by the format, or switch the saved format to an unsafe-free serialization; passingmap_locationtotorch.loadis also clearer.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/synference/simformer/model.py` around lines 538 - 578, Harden SimformerModel.load by replacing unsafe full-pickle loading with weights_only=True and registering the NumPy dtype/array globals required by the saved payload, while passing the method’s map_location argument directly to torch.load. Preserve the existing format-version validation and reconstruction flow.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/synference/simformer/masks.py`:
- Around line 57-69: Update random_condition_mask so the Beta probability is
derived from generator-controlled Gamma draws rather than beta_dist.sample(),
which uses the global RNG. Construct the Beta-equivalent ratio using alpha and
beta parameters, pass the supplied generator to the underlying draws, and
preserve the existing mask shape and _fix_all_true_rows behavior.
In `@src/synference/simformer/model.py`:
- Around line 152-160: The _as_condition_mask method currently returns CPU masks
that are later used to index device-local tensors. Move the validated mask to
self.device before returning it, and update sample_intervals so constraint_mask
and effective_condition conversions follow the same device placement; ensure
_init_x_T, _prepare_x_o, _extract_latents, and log_prob can index without
cross-device errors.
In `@src/synference/simformer/train.py`:
- Around line 374-394: Update the final checkpoint restoration condition around
best_state so the model loads the best validation checkpoint whenever best_state
is available, regardless of whether stats["early_stopped"] is true. Preserve the
existing guard that avoids loading when no validation checkpoint was recorded.
- Around line 277-281: Shuffle the input rows with a seeded permutation before
computing the validation split in train_simformer, so data_val is sampled
independently of caller ordering while preserving reproducibility. Apply the
same permuted tensor to both data_val and data_train, and retain the existing
validation repeat and empty-validation behavior.
---
Nitpick comments:
In `@src/synference/sbi_runner.py`:
- Around line 8397-8402: In sample_posterior, replace the assert guarding the
X_test column count against mask.sum() with explicit user-facing validation that
raises an appropriate exception containing the existing mismatch details.
Preserve the current shape normalization and allow valid column/mask matches to
proceed unchanged.
In `@src/synference/simformer/model.py`:
- Around line 538-578: Harden SimformerModel.load by replacing unsafe
full-pickle loading with weights_only=True and registering the NumPy dtype/array
globals required by the saved payload, while passing the method’s map_location
argument directly to torch.load. Preserve the existing format-version validation
and reconstruction flow.
In `@src/synference/simformer/sampling.py`:
- Around line 179-203: Address the scalability of exact divergence in
_divergence_exact by either switching to Hutchinson divergence above a clearly
defined latent-dimension threshold or documenting the num_steps × n_latent
backward-pass cost at the public SimformerModel.log_prob and
Simformer_Fitter.log_prob entry points. Preserve exact divergence behavior below
any threshold and keep the existing log-probability API consistent.
In `@tests/test_simformer.py`:
- Around line 405-427: Mark the model-training test fixtures or their containing
test classes, including the two-moons fixture two_moons_model and the
galaxy-library fixture covering the other referenced range, with
pytest.mark.slow so they are excluded from the default fast suite while
remaining available through CI’s slow-test opt-in.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 567b7286-faa7-4b66-8679-be3d1d2d15e8
📒 Files selected for processing (18)
CLAUDE.mdREADME.mddocs/source/advanced_topics/simformer.ipynbdocs/source/getting_started/installation.rstexamples/sbi/scripts/train_model.pyexamples/simformer/scripts/fit_simformer_model.pyexamples/simformer/scripts/two_moons_demo.pypyproject.tomlsrc/synference/__init__.pysrc/synference/sbi_runner.pysrc/synference/simformer/__init__.pysrc/synference/simformer/masks.pysrc/synference/simformer/model.pysrc/synference/simformer/nn.pysrc/synference/simformer/sampling.pysrc/synference/simformer/sde.pysrc/synference/simformer/train.pytests/test_simformer.py
💤 Files with no reviewable changes (4)
- pyproject.toml
- README.md
- docs/source/getting_started/installation.rst
- examples/sbi/scripts/train_model.py
| def random_condition_mask( | ||
| num_samples: int, | ||
| theta_dim: int, | ||
| x_dim: int, | ||
| generator: Optional[torch.Generator] = None, | ||
| alpha: float = 1.0, | ||
| beta: float = 4.0, | ||
| ) -> torch.Tensor: | ||
| """Bernoulli masks with per-batch probability drawn from ``Beta(alpha, beta)``.""" | ||
| beta_dist = torch.distributions.Beta(alpha, beta) | ||
| prob = beta_dist.sample() | ||
| mask = torch.rand(num_samples, theta_dim + x_dim, generator=generator) < prob | ||
| return _fix_all_true_rows(mask) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Does torch.distributions.Beta.sample accept a generator argument?
💡 Result:
As of July 25, 2026, the torch.distributions.Beta.sample method does not natively accept a generator argument [1][2][3]. While some updates to the PyTorch torch.distributions library have been initiated to standardize generator support across various distributions [4][5], the sample and rsample methods for the Beta distribution currently maintain their standard signature, which typically accepts only sample_shape [1][2]. Users requiring control over the random number generation process for reproducibility generally rely on setting the global PyTorch manual seed (torch.manual_seed()) before calling the sample() method [6]. Recent development activity in the PyTorch repository (as of March 2026) indicates ongoing efforts to add a generator parameter to the base distribution classes [5], but users should verify the specific version of PyTorch they are using to determine if this has been implemented for the Beta class. To check your current environment's capabilities, you can inspect the signature of your distribution instance's sample method using Python's inspect module.
Citations:
- 1: https://github.com/pytorch/pytorch/blob/v2.12.0/torch/distributions/beta.py
- 2: https://docs.pytorch.org/docs/2.3/%5Fmodules/torch/distributions/beta.html
- 3: https://docs.pytorch.org/docs/stable/distributions.md
- 4: Standardized Distributions pytorch/pytorch#45115
- 5: Fix: Add generator parameter support to Gamma distribution sample() pytorch/pytorch#177448
- 6: https://stackoverflow.com/questions/67901636/taking-sample-from-categorical-distribution-pytorch
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate and inspect the relevant files/interfaces without executing repository code.
printf 'Tracked python files matching paths:\n'
git ls-files | rg '(^|/)masks\.py$|(^|/)train_simformer\.py$|(^|/)synference\.py$' || true
printf '\nRelevant definitions/usages:\n'
rg -n "def random_condition_mask|random_condition_mask\(|torch\.manual_seed|random_generatr|generator|Beta\\(" -S .
printf '\nTarget masked section:\n'
sed -n '1,120p' "$(git ls-files | rg '/masks\.py$' | head -n 1)"
printf '\nTrain caller references:\n'
for f in $(git ls-files | rg 'train_simformer\.py$|simformer.*\.py$'); do
if rg -q "random_condition_mask" "$f"; then
echo "--- $f"
rg -n -C 5 "random_condition_mask|torch\.manual_seed|random_generatr|Generator" "$f"
fi
doneRepository: synthesizer-project/synference
Length of output: 24253
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the Simformer training path around the generator seed and mask call.
sed -n '235,330p' src/synference/simformer/train.py
printf '\nConfiguration/name references for random condition mask:\n'
rg -n --hidden -S '"random"|random_condition_mask|condition_mask.*random|condition_mask_fn|random"%' . || trueRepository: synthesizer-project/synference
Length of output: 7343
Generate the Beta probability with the passed RNG.
random_condition_mask re-seeds generator inside the call, but torch.distributions.Beta(...).sample() does not accept a generator argument, so its global RNG still determines prob. Use a seeded Beta proxy, e.g. Gamma draws, to keep this sampler fully reproducible.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/synference/simformer/masks.py` around lines 57 - 69, Update
random_condition_mask so the Beta probability is derived from
generator-controlled Gamma draws rather than beta_dist.sample(), which uses the
global RNG. Construct the Beta-equivalent ratio using alpha and beta parameters,
pass the supplied generator to the underlying draws, and preserve the existing
mask shape and _fix_all_true_rows behavior.
| def _as_condition_mask(self, condition_mask) -> torch.Tensor: | ||
| mask = torch.as_tensor(np.asarray(condition_mask), dtype=torch.bool) | ||
| if mask.shape != (self.num_nodes,): | ||
| raise ValueError( | ||
| f"condition_mask must have shape ({self.num_nodes},), got {tuple(mask.shape)}." | ||
| ) | ||
| if mask.all(): | ||
| raise ValueError("condition_mask cannot condition on every node.") | ||
| return mask |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check whether any caller/test exercises the Simformer model on a non-CPU device.
rg -nP --type=py -C3 '(to\(\s*["'\'']cuda|device\s*=\s*["'\'']cuda|SimformerModel[^\n]*\.to\()' src tests examplesRepository: synthesizer-project/synference
Length of output: 529
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Files matching model.py:"
fd -a 'model\.py$' . || true
echo
echo "Inspect src/synference/simformer/model.py around relevant symbols:"
wc -l src/synference/simformer/model.py
sed -n '130,310p' src/synference/simformer/model.py
sed -n '460,510p' src/synference/simformer/model.py
sed -n '380,415p' src/synference/simformer/model.pyRepository: synthesizer-project/synference
Length of output: 12743
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '400,525p' src/synference/simformer/model.pyRepository: synthesizer-project/synference
Length of output: 5773
🏁 Script executed:
#!/bin/bash
set -euo pipefail
if command -v python3 >/dev/null 2>&1; then
python3 - <<'PY'
import subprocess, tempfile, os, textwrap
tmpdir = tempfile.mkdtemp()
src = os.path.join(tmpdir, "test_device_mask.py")
with open(src, "w") as f:
f.write(textwrap.dedent('''
import torch
x = torch.arange(2 * 3.0).reshape(2, 3).cuda() if torch.cuda.is_available() else torch.arange(2 * 3.0).reshape(2, 3)
cpu_mask = torch.tensor([True, False, True])
try:
result = x[:, cpu_mask]
print("cpu_mask_indexing_succeeded")
except Exception as e:
print(type(e).__name__)
print(str(e).splitlines()[0])
if torch.cuda.is_available():
gpu_mask = cpu_mask.cuda()
result = x[:, gpu_mask]
print("gpu_mask_indexing_succeeded")
'''))
subprocess.run([sys.executable, src], text=True)
PY
else
echo "python3 not available"
fiRepository: synthesizer-project/synference
Length of output: 279
🌐 Web query:
PyTorch RuntimeError expected index Tensor to be on the same device as the base tensor boolean mask device mismatch indexing CUDA tensor CPU tensor
💡 Result:
The PyTorch RuntimeError "expected index Tensor to be on the same device as the base tensor" (or similar device mismatch errors) occurs because PyTorch operations, including indexing and masking, require all input tensors to reside on the same device [1][2][3]. Unlike some other operations, PyTorch does not automatically transfer tensors between devices (e.g., from CPU to GPU) during indexing to prevent unexpected performance bottlenecks and silent data movement [1]. To resolve this issue, you must ensure that your index tensor (the boolean mask or integer index) and the base tensor (the data being indexed) are explicitly placed on the same device [1][3]. How to fix it: 1. Identify the device of your tensors: You can check the device of any tensor by inspecting its.device attribute [3]. print(base_tensor.device) print(mask_tensor.device) 2. Move the index/mask tensor to the target device: If your base tensor is on the GPU (e.g., 'cuda:0') and your mask is on the 'cpu', move the mask to the GPU using the.to method [1][4][5]. mask_tensor = mask_tensor.to(base_tensor.device) # Now perform the indexing result = base_tensor[mask_tensor] Common scenarios for this error: - Data loaders: Tensors created in a custom Dataset or during data preprocessing often default to the CPU. If you move your model to the GPU but forget to move these auxiliary tensors, the error will trigger [1]. - Boolean masks: Masks created using conditions (e.g., mask = tensor > 0) inherit the device of the original tensor. However, if you create a mask manually (e.g., using torch.ones or torch.tensor) without specifying the device, it will default to 'cpu' [5]. Always specify the device when creating new tensors intended for GPU operations [5]: mask = torch.ones(size, device=base_tensor.device, dtype=torch.bool) - Library functions: Some functions, such as those inside nn.functional or third-party modules, might internally create masks or buffers on the CPU [5][6]. Ensure any tensors passed to these functions are on the same device as the function's expected inputs [2][5]. If you are using a specific device like MPS (Apple Silicon), ensure you are on a recent PyTorch version, as older versions had limited support for cross-device indexing on non-CUDA backends [7].
Citations:
- 1: https://adhdecode.com/debugging/pytorch/runtimeerror-indices-should-be-either-on-cpu-or-on/
- 2: https://discuss.pytorch.org/t/runtimeerror-expected-all-tensors-to-be-on-the-same-device-but-found-at-least-two-devices-cpu-and-cuda-0-when-checking-arugment-for-argument-index-in-method-wrapper-index-select/138437
- 3: https://discuss.pytorch.org/t/any-way-to-improve-device-mismatch-errors/202950
- 4: https://stackoverflow.com/questions/74980419/how-to-find-what-is-using-the-cpu-for-indices-pytorch-numpy-indices-should-b
- 5: Device mismatch in scaled_dot_product_attention docstring example pytorch/pytorch#166117
- 6: Mask in MaskedTensor does not change device pytorch/pytorch#147140
- 7: MPS runtime error when indexing from list / tensor pytorch/pytorch#81051
Move condition masks to self.device before indexing.
_as_condition_mask() creates boolean masks on CPU, but _init_x_T, _prepare_x_o, _extract_latents, and log_prob index tensors allocated on self.device. That causes cross-device indexing failures on CUDA; move the masks in _as_condition_mask() and handle the sample_intervals constraint_mask/effective_condition conversion consistently.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/synference/simformer/model.py` around lines 152 - 160, The
_as_condition_mask method currently returns CPU masks that are later used to
index device-local tensors. Move the validated mask to self.device before
returning it, and update sample_intervals so constraint_mask and
effective_condition conversions follow the same device placement; ensure
_init_x_T, _prepare_x_o, _extract_latents, and log_prob can index without
cross-device errors.
| # Validation split (first fraction of rows, as in the original implementation). | ||
| n_val = max(int(train_config["validation_fraction"] * data.shape[0]), 0) | ||
| data_t = torch.as_tensor(data, device=device) | ||
| data_val = data_t[:n_val].repeat(train_config["val_repeat"], 1) if n_val > 0 else None | ||
| data_train = data_t[n_val:] |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Validation split takes the leading rows without shuffling.
data_val is the first validation_fraction of rows as supplied. Callers that pass a library in generation order (e.g. sorted by a parameter) get a systematically biased validation set, which then drives early stopping. The Simformer_Fitter path happens to shuffle upstream, but direct train_simformer users are exposed.
A seeded permutation before splitting would remove the dependency on caller ordering:
🛡️ Proposed fix
n_val = max(int(train_config["validation_fraction"] * data.shape[0]), 0)
+ perm = torch.randperm(data.shape[0], generator=generator).numpy()
+ data = data[perm]
data_t = torch.as_tensor(data, device=device)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Validation split (first fraction of rows, as in the original implementation). | |
| n_val = max(int(train_config["validation_fraction"] * data.shape[0]), 0) | |
| data_t = torch.as_tensor(data, device=device) | |
| data_val = data_t[:n_val].repeat(train_config["val_repeat"], 1) if n_val > 0 else None | |
| data_train = data_t[n_val:] | |
| # Validation split (first fraction of rows, as in the original implementation). | |
| n_val = max(int(train_config["validation_fraction"] * data.shape[0]), 0) | |
| perm = torch.randperm(data.shape[0], generator=generator).numpy() | |
| data = data[perm] | |
| data_t = torch.as_tensor(data, device=device) | |
| data_val = data_t[:n_val].repeat(train_config["val_repeat"], 1) if n_val > 0 else None | |
| data_train = data_t[n_val:] |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/synference/simformer/train.py` around lines 277 - 281, Shuffle the input
rows with a seeded permutation before computing the validation split in
train_simformer, so data_val is sampled independently of caller ordering while
preserving reproducibility. Apply the same permuted tensor to both data_val and
data_train, and retain the existing validation repeat and empty-validation
behavior.
| if val_loss < min_val_loss: | ||
| min_val_loss = val_loss | ||
| best_state = copy.deepcopy({k: v.cpu() for k, v in net.state_dict().items()}) | ||
| if early_stopping_counter > train_config["stop_early_count"]: | ||
| stats["early_stopped"] = True | ||
| if verbose: | ||
| logger.info(f"Early stopping at step {step} (val loss {val_loss:.4f}).") | ||
| break | ||
|
|
||
| if verbose and (step % print_every) == 0: | ||
| message = f"Step {step}/{total_steps}: train loss {train_loss_ema:.4f}" | ||
| if stats["val_loss"]: | ||
| message += f", val loss {stats['val_loss'][-1]:.4f}" | ||
| logger.info(message) | ||
| if progress_callback is not None: | ||
| progress_callback( | ||
| step, train_loss_ema, stats["val_loss"][-1] if stats["val_loss"] else None | ||
| ) | ||
|
|
||
| if stats["early_stopped"] and best_state is not None: | ||
| net.load_state_dict(best_state) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Best checkpoint is discarded when training runs to completion.
best_state is tracked on every validation improvement but only restored under stats["early_stopped"]. A full-budget run therefore returns the final weights even when a clearly better validation point was seen earlier. Restoring best_state whenever it exists (or documenting the deliberate choice) would make the two paths consistent.
♻️ Proposed change
- if stats["early_stopped"] and best_state is not None:
+ if best_state is not None:
net.load_state_dict(best_state)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if val_loss < min_val_loss: | |
| min_val_loss = val_loss | |
| best_state = copy.deepcopy({k: v.cpu() for k, v in net.state_dict().items()}) | |
| if early_stopping_counter > train_config["stop_early_count"]: | |
| stats["early_stopped"] = True | |
| if verbose: | |
| logger.info(f"Early stopping at step {step} (val loss {val_loss:.4f}).") | |
| break | |
| if verbose and (step % print_every) == 0: | |
| message = f"Step {step}/{total_steps}: train loss {train_loss_ema:.4f}" | |
| if stats["val_loss"]: | |
| message += f", val loss {stats['val_loss'][-1]:.4f}" | |
| logger.info(message) | |
| if progress_callback is not None: | |
| progress_callback( | |
| step, train_loss_ema, stats["val_loss"][-1] if stats["val_loss"] else None | |
| ) | |
| if stats["early_stopped"] and best_state is not None: | |
| net.load_state_dict(best_state) | |
| if val_loss < min_val_loss: | |
| min_val_loss = val_loss | |
| best_state = copy.deepcopy({k: v.cpu() for k, v in net.state_dict().items()}) | |
| if early_stopping_counter > train_config["stop_early_count"]: | |
| stats["early_stopped"] = True | |
| if verbose: | |
| logger.info(f"Early stopping at step {step} (val loss {val_loss:.4f}).") | |
| break | |
| if verbose and (step % print_every) == 0: | |
| message = f"Step {step}/{total_steps}: train loss {train_loss_ema:.4f}" | |
| if stats["val_loss"]: | |
| message += f", val loss {stats['val_loss'][-1]:.4f}" | |
| logger.info(message) | |
| if progress_callback is not None: | |
| progress_callback( | |
| step, train_loss_ema, stats["val_loss"][-1] if stats["val_loss"] else None | |
| ) | |
| if best_state is not None: | |
| net.load_state_dict(best_state) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/synference/simformer/train.py` around lines 374 - 394, Update the final
checkpoint restoration condition around best_state so the model loads the best
validation checkpoint whenever best_state is available, regardless of whether
stats["early_stopped"] is true. Preserve the existing guard that avoids loading
when no validation checkpoint was recorded.
1 participant
Written in torch rather than jax, this reimplements the Simformer approach to learn full conditional distributions.
Summary by CodeRabbit
New Features
Documentation
Chores