mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-25 15:56:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
25,761 Commits 508 Branches 263 Tags
Commit Graph

7 Commits

Author SHA1 Message Date
Matthew Johnson
719031c1fd [mutable-arrays] persist shardings through xla computations 2025-02-07 18:33:24 +00:00
Matthew Johnson
b6482f126e add mutable array ref error checks to cond and custom_vjp 2024-12-20 01:44:50 +00:00
Matthew Johnson
e52856261f add mutable array ref error checks to scan 2024-12-19 01:33:39 +00:00
Matthew Johnson
42ac4ca357 ref errors 2024-12-18 07:46:14 +00:00
Dougal
fc2edbfac8 Add a freeze primitive to delimit ref lifetimes for AD. 
Also some basic AD through mutable_array/freeze.

Co-authored-by: Matthew Johnson <mattjj@google.com>
 2024-12-09 20:57:07 -05:00
Sharad Vikram
80f963c003 Fix mutable array effects not being tracked properly 
PiperOrigin-RevId: 680801564
 2024-09-30 18:55:15 -07:00
Dougal
e63b35d550 Add discharge rules for scan with mutable arrays. Move mutable array tests to separate file. 
Co-authored-by: Matt Johnson <mattjj@google.com>
 2024-05-02 14:36:16 -04:00