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

34 Commits

Author SHA1 Message Date
Michael Hudgins
d4d1518c3d Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax 
PiperOrigin-RevId: 676843138
 2024-09-20 07:52:33 -07:00
Jake VanderPlas
f090074d86 Avoid 'from jax import config' imports 
In some environments this appears to import the config module rather than
the config object.
 2024-04-11 13:23:27 -07:00
Jake VanderPlas
4edb74ba7b Fix some numpy 2.0 incompatibilities 2023-09-21 10:24:52 -07:00
Jake VanderPlas
2f878a7168 Tests: set jax_legacy_prng_key='error' 2023-08-28 10:56:09 -07:00
Jake VanderPlas
fbe4f10403 Change to simpler import for jax.config 2023-04-21 11:51:22 -07:00
Jake VanderPlas
8431e43fe5 [x64] more type safety in stax_test.py 2022-12-02 10:02:25 -08:00
Peter Hawkins
72f4f389be Migrate remaining tests from jtu.cases_from_list to jtu.sample_product. 
Delete jtu.cases_from_list.
 2022-10-12 15:20:53 +00:00
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
Jeppe Klitgaard
17de89b16a feat: refactor code using pyupgrade 
This PR upgrades legacy Python code to 3.7+ code using pyupgrade:
```sh
pyupgrade --py37-plus --keep-runtime-typing **.py
```

a
 2022-05-17 22:14:05 +01:00
Jake VanderPlas
e376df29be disable implicit rank promotion in a number of remaining tests 2022-01-28 08:16:30 -08:00
Jake VanderPlas
df0969961b Testing: avoid hard-coding random seeds 2021-12-10 10:32:09 -08: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
Peter Hawkins
db2e91eba2 Move jax.test_util to jax._src.test_util. 
Add forwarding shims for names used by external clients of JAX in practice.

PiperOrigin-RevId: 398721725
 2021-09-24 07:02:49 -07:00
Roman Novak
b65f39ca7a Default to jnp.float_ type in nn.initializers. 2021-08-17 20:41:09 -07:00
Jake VanderPlas
afce718eb1 Add ability to specify individual test targets 2020-06-29 11:08:57 -07:00
joao guilherme
77e4d8b3b9
Updates onp -> np in random, loops, jet and in the tests of stax and optix (#3182) 2020-05-21 14:12:18 -07:00
Madison May
4fa153439e FIX: batch norm w/ no scale, center 2020-03-31 19:59:57 -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
Roman Novak
ec4b544508 Stax: allow arbitrary dimension numbers in pooling layers. (#1892) 2019-12-18 15:27:46 -08:00
fehiepsi
5ffe2ae5dd expose sigmoid too 2019-08-01 17:11:31 -04:00
fehiepsi
688d77f432 better test str message 2019-08-01 16:41:06 -04:00
fehiepsi
2836d03b5e add some nonlinearity 2019-08-01 16:39:08 -04:00
Anselm Levskaya
aee66b5a24
comment to trigger travis rebuild 2019-04-10 00:57:02 -07:00
Anselm Levskaya
cef4c94c13 finish transposed convolution implementation and tests 2019-04-09 22:59:03 -07:00
Jamie Townsend
1c9b9a57fd Use jax.random for stax initialization 2019-04-03 12:54:02 +01:00
Peter Hawkins
d2a2112303 Support non-default axis values in batch norm. 
Fixes #461.
 2019-03-14 09:32:53 -04:00
Anselm Levskaya
c3cbf4a20a allow apply_fns to take general kwargs, not just rng 2019-03-01 14:24:50 -08:00
Peter Hawkins
274a63e518 Fix implementation of average pooling to align the window element counts with the spatial dimensions. 2019-01-28 09:20:02 -05:00
Matthew Johnson
ecaae6bdd0 add Softmax layer to stax (closes #182) 2019-01-05 10:06:31 -08:00
Matthew Johnson
b4246163ac add stax.FanInConcat (fixes #174) 2018-12-30 16:51:32 -08:00
Matthew Johnson
d3ec0b23c5 fix miscellaneous bugs, incl. complex abs grad 2018-12-12 12:03:57 -05:00
Dougal Maclaurin
8b88027df0 Number of test cases settable with command-line flag 2018-12-06 18:30:59 -05:00
Dougal Maclaurin
2df36f7510 Made a shim to handle configuration without having absl parse command-line flags. 
PiperOrigin-RevId: 223391288
 2018-11-29 13:44:54 -08:00
Matthew Johnson
489dd1c81a move jax tests and examples to repo root 2018-11-18 15:15:47 -08:00