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

12 Commits

Author SHA1 Message Date
Peter Hawkins
ba557d5e1b Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.". 
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.

PiperOrigin-RevId: 476167538
 2022-09-22 12:27:19 -07:00
Roy Frostig
623c201054 [JAX] move example libraries from jax.experimental into jax.example_libraries 
The `jax.experimental.stax` and `jax.experimental.optimizers` modules are standalone examples libraries. By contrast, the remaining modules in `jax.experimental` are experimental features of the JAX core system. This change moves the two example libraries, and the README that describes them, to `jax.example_libraries` to reflect this distinction.

PiperOrigin-RevId: 404405186
 2021-10-19 17:30:45 -07:00
Skye Wanderman-Milne
cd619978bb Adjust precision in examples/kernel_lsq.py and corresponding test. 
This is important on TPU.
 2021-03-11 03:03:43 +00:00
Jake Vanderplas
6aa8f2461c
Fix remaining flakes and use exclude within setup.cfg (#3371) 2020-06-08 22:58:03 -07:00
Peter Hawkins
d59ecddfe8
Replace np -> jnp, onp -> np in examples/ (#2971) 
For context, see #2370
 2020-05-05 15:45:07 -04:00
Peter Hawkins
e60d5dd54c
Remove "from __future__" uses from JAX. (#2117) 
The future (Python 3) has arrived; no need to request it explicitly.
 2020-01-29 12:29:03 -05:00
Peter Hawkins
dcc882cf6b
Drop Python 2 support from JAX. (#1962) 
Remove six dependency.
 2020-01-08 13:17:55 -05:00
Matthew Johnson
642d2dc802 revies optimizers api, fix misc bugs 
* add more optimizers numerical tests
* update examples and readme with new optimziers api
* add device_values parameter to xla_call
* change optimizers.py to flatten trees and subtrees
* remove tree_map2, tree_multimap2, tree_mimomap, tree_prefixmap
* add optimizer tests: DeviceTuples and error msgs
* make the device_values arg to jit private
 2019-05-03 12:44:52 -07:00
Matthew Johnson
5847d96168 rename "minmax" -> "optimizers" 2019-02-06 11:02:16 -08:00
Roy Frostig
dda5bdb68b illustrate two-level autobatching with jaxprs of the gram matrix function in the kernel least-squares example 2018-12-16 13:27:36 -08:00
Peter Hawkins
0d4eb6c1e1 Make JAX flake8-clean. 
Fixes #1.
 2018-12-13 15:29:39 -05:00
Roy Frostig
693e51073e add kernel least-squares regression example 2018-12-12 17:22:33 -08:00