Skip to content

Pull requests: google-deepmind/torax

New pull request New
47 Open 2,072 Closed
47 Open 2,072 Closed
Author
Filter by author
Loading
Label
Filter by label
Loading
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Loading
Milestones
Filter by milestone
Loading
Reviews
Filter by reviews
No reviews Review required Approved review Changes requested
Assignee
Filter by who’s assigned
Assigned to nobody Loading
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

Use jax.closure_convert to make jax_utils.while_loop_bounded with implementation='while_loop' be differentiable through closures.
#2303 opened Jul 23, 2026 by copybara-service Bot Loading…
Add JVP support to jax_utils.while_loop_bounded using the new JAX hijax API, which allows both custom VJP and custom JVP.
#2302 opened Jul 22, 2026 by copybara-service Bot Loading…
Add JVP test for jax_utils.while_loop_bounded against jax.lax.while_loop`.
#2301 opened Jul 22, 2026 by copybara-service Bot Loading…
Add explicit radial smoothing zones to CombinedTransportModel.
#2300 opened Jul 22, 2026 by copybara-service Bot Loading…
Optimize unstacking of JIT outputs.
#2296 opened Jul 20, 2026 by copybara-service Bot Loading…
Switch while_loop_bounded to use implementation='while_loop'.
#2295 opened Jul 20, 2026 by copybara-service Bot Loading…
Minor plot improvements copybara:import-manual Set when ready for copybara manual import
#2291 opened Jul 16, 2026 by theo-brown Collaborator Loading…
1
Add test for forward-mode differentiation.
#2289 opened Jul 15, 2026 by copybara-service Bot Loading…
Add config/CLI option to select a single slider mode
#2286 opened Jul 14, 2026 by Neal006 Loading…
2
Reduce plotting slider payload for faster figure load
#2285 opened Jul 14, 2026 by Neal006 Loading…
2
Support different sources for trapped particle fraction
#2283 opened Jul 13, 2026 by theo-brown Collaborator Loading…
14
Migrate ITER Predictor-Corrector configs to combined transport model.
#2260 opened Jun 29, 2026 by copybara-service Bot Loading…
Migrate default and simple configs to combined transport model.
#2256 opened Jun 29, 2026 by copybara-service Bot Loading…
Migrate SPARC configs to combined transport model.
#2254 opened Jun 29, 2026 by copybara-service Bot Loading…
Migrate ITER configs to combined transport model.
#2253 opened Jun 29, 2026 by copybara-service Bot Loading…
Code update
#2252 opened Jun 29, 2026 by copybara-service Bot Loading…
Followup PR: Implement IMAS output core_sources mapping
#2231 opened Jun 18, 2026 by MateoBell Collaborator Loading…
(PR 2/2) Add main_ions contribution to output core_profiles mapping
#2227 opened Jun 18, 2026 by MateoBell Collaborator Loading…
1
2
Add mapping for output core_transport to IMAS
#2224 opened Jun 18, 2026 by MateoBell Collaborator Loading…
9
(PR 4/4): Implement IMAS output core_sources mapping
#2222 opened Jun 17, 2026 by MateoBell Collaborator Loading…
Add normalized collisionality nu_star to outputs
#2203 opened Jun 9, 2026 by r1cksync Loading…
1
2
fix: use left-biased 3-point stencil for face gradients on non-uniform grids
#2202 opened Jun 9, 2026 by copybara-service Bot Loading…
Add gyaradax plugin
#2189 opened Jun 3, 2026 by gerkone • Draft
3 of 7 tasks
2
Add explicit data dictionary version dependency.
#2155 opened May 6, 2026 by copybara-service Bot Loading…
2
Expose Anderson acceleration options in the config.
#2154 opened May 6, 2026 by copybara-service Bot Loading…
ProTip! Updated in the last three days: updated:>2026-07-20.