Skip to content

ENH(zuko): restore zuko training controls - #14

Open
mj-will wants to merge 4 commits into
mainfrom
restore-zuko-training-controls
Open

mj-will wants to merge 4 commits into
mainfrom
restore-zuko-training-controls

Conversation

@mj-will

@mj-will mj-will commented Oct 31, 2026 •

Copy link
Copy Markdown
Contributor

Add the training controls present in aspire for the zuko flows.

Helps to address mj-will/aspire#89

Add gradient norm clipping, cosine learning-rate annealing, patience and opt-in progress, preserving best-epoch restoration and complete histories. Validate raw and transformed training particles before splitting.

Cover numerical optimization, RNG isolation and Sequince fit_kwargs integration. Document defaults and validation_fraction=0 behavior.

Validation: 42 Zuko tests and 93 full-suite tests pass on CPU; all pre-commit hooks and the strict Pixi docs build pass. GPU testing unavailable.

Refs: mj-will/aspire#89
Extract the Zuko validation helper into density utilities using the native array namespace. Preserve shape checks and dimension-specific errors, and cover NumPy, Torch and JAX inputs.

Validation: 96 full-suite tests passed; all pre-commit checks and the strict Pixi docs build passed. GPU testing unavailable.
@read-the-docs-community

read-the-docs-community Bot commented Oct 31, 2026 •

Copy link
Copy Markdown

Documentation build overview

📚 sequince-density | 🛠️ Build #35542896 | 📁 Comparing c717916 against latest (bb64e71)

  🔍 Preview build  

5 files changed · ± 5 modified

± Modified

Keep final parameters and an empty validation history when validation is disabled. Reject patience in that mode; retain best-epoch selection when validation is enabled.

Validation: 50 Zuko tests and 104 full-suite tests passed on CPU, with no skips. All pre-commit hooks and the strict docs build passed. GPU testing unavailable.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant