mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-05-02 08:16:08 +00:00
Code Issues Packages Projects Releases Wiki Activity
8,838 Commits 508 Branches 263 Tags
Commit Graph

3 Commits

Author SHA1 Message Date
Sergei Lebedev
af41a959d3 Most of JAX now uses concrete types for things defined in jaxlib.xla_client 
Note that a few call sites in the diff got a ``# type: ignore``, because
the latest jaxlib does not have up-to-date signatures for the correpsonding
callables.
 2021-08-16 20:33:36 +01:00
Roy Frostig
b8f9dd6269 unify tree_map and tree_multimap 2021-04-28 19:59:31 -07:00
Peter Hawkins
6ee6c59235 Move jax.tree_util implementation to jax._src.tree_util. 
NFC intended.

PiperOrigin-RevId: 364857920
 2021-03-24 12:00:38 -07:00