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

25 Commits

Author SHA1 Message Date
Matthew Johnson
c7ce442084 initial pmap/pxla code, pair-coded w/ @dougalm 2019-01-24 13:28:54 -08:00
Matthew Johnson
42bd253cbb add prng design doc in markdown format 2019-01-18 07:38:43 -08:00
Peter Hawkins
86d8915c3d Add Sphinx-generated reference documentation for JAX. 2019-01-16 09:13:31 -05:00
Matthew Johnson
54886bd310 add sort_key_val batching rule (fixes #221) 2019-01-13 11:56:36 -08:00
Matthew Johnson
c3664e95e9 set 'jacobian' to point to jacrev 2019-01-07 08:56:19 -08:00
Matthew Johnson
0f7c7c4eab generalize jacfwd and jacrev to handle pytrees 2019-01-06 12:49:41 -08:00
Matthew Johnson
ad4322c5da playing around with flattening functions 2019-01-06 12:49:35 -08:00
Matthew Johnson
8f1bc997ca add value-and-grad fun (closes #149) 2018-12-20 10:09:34 -08:00
Dougal Maclaurin
87922fdf13 Generalized make_jaxpr to handle python containers 2018-12-19 10:59:13 -05:00
Roy Frostig
e0ad5bb394 add TODO comment for containers in api.make_jaxpr 2018-12-17 11:42:45 -08:00
Roy Frostig
f4a8e03ce1 add a basic make_jaxpr transformation to the api module 2018-12-16 13:26:02 -08:00
Roy Frostig
3b8fdb050a wrap "jit" around generated function name 2018-12-16 13:24:20 -08:00
Roy Frostig
b318f56928 generate name, module, and doctring for functions output from jit. 2018-12-16 13:04:29 -08:00
Dougal Maclaurin
9b835a180f Added some docstrings to top-level transformations 2018-12-12 09:47:49 -05:00
Matthew Johnson
89349e5e6d fix transpose issue in jacfwd and jacrev 2018-12-11 16:24:20 -08:00
Matthew Johnson
cbd9ace425 change vmap api to be curried (closes #78) 2018-12-11 12:52:09 -08:00
Matthew Johnson
b907bcd5ae improve jaxlib wheel building script 2018-12-09 08:59:08 -08:00
Dougal Maclaurin
d89c7c2457 Merge branch 'master' into quickish-check 2018-12-07 11:13:41 -05:00
Matthew Johnson
50624bd978 rename in_bdims, out_bdims --> in_axes, out_axes 2018-12-07 06:53:29 -08:00
Dougal Maclaurin
2e4ff400fa Fixed bug in vjp with constant-zero tangent outputs 2018-12-05 15:55:01 -05:00
Dougal Maclaurin
307d195577 Wrapped static args to jit to be hashed on id. This is conservative but simple and cheap. 2018-12-05 15:55:01 -05:00
Matthew Johnson
9101d66b4e source sync 
PiperOrigin-RevId: 222487460
 2018-11-21 20:32:33 -08:00
Matthew Johnson
3b3490f406 source sync 
PiperOrigin-RevId: 222483357
 2018-11-21 20:32:33 -08:00
Peter Hawkins
5e60639bc5 source sync 
PiperOrigin-RevId: 222452709
 2018-11-21 20:22:54 -08:00
Matthew Johnson
a30e858e59 populating source tree 2018-11-17 18:03:33 -08:00