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
