92 Commits

Author SHA1 Message Date
Peter Hawkins
1b3dd65daf
Avoid lexically capturing the train_images value in MNIST VAE example. (#3947)
* Avoid lexically capturing the train_images value in MNIST VAE example.

This has the effect of baking in the training dataset as a constant, something that LLVM does not like that much.

* Add device_put to images.
2020-08-03 16:50:20 -04:00
Matthew Johnson
d10cf0e38f
fix prng key reuse in differential privacy example (#3646)
fix prng key reuse in differential privacy example
2020-07-02 14:29:17 -07:00
Jake VanderPlas
afce718eb1 Add ability to specify individual test targets 2020-06-29 11:08:57 -07:00
Jake Vanderplas
6aa8f2461c
Fix remaining flakes and use exclude within setup.cfg (#3371) 2020-06-08 22:58:03 -07:00
Matthew Johnson
c42a7f7890
remove some trailing whitespace (#3287) 2020-06-02 17:37:20 -07:00
Peter Hawkins
bc5a0b336b
Remove some uses of jax.partial. (#3131) 2020-05-18 10:19:03 -04:00
Peter Hawkins
d59ecddfe8
Replace np -> jnp, onp -> np in examples/ (#2971)
For context, see #2370
2020-05-05 15:45:07 -04:00
Jin Dong
60d856ab9f remove from __future__ code 2020-04-09 18:16:47 -04:00
Xiayun Sun
651316f4c7
Fix issue 1465: fix jit in example (#1473)
* fix jit in example

* Avoid using static_argnums on a keyword argument; use a positional argument and a wrapper function for now.

Co-authored-by: Peter Hawkins <phawkins@google.com>
2020-04-01 11:57:57 -04:00
Matthew Johnson
b015e57169 try re-enabling control tests that trigger #2507 2020-03-30 20:12:33 -07:00
Matthew Johnson
7e480fa923 add custom_jvp / vjp, delete custom_transforms 2020-03-21 22:08:03 -07:00
Srinivas Vasudevan
c7f211d433
Update JAX to use XLA hyperbolic functions. (#2415) 2020-03-19 10:29:37 -04:00
bddppq
ac6a313cfc
Fix ONNX mnist example (#2374)
* Fix ONNX mnist example

* use np to compute the shape; rename jax.numpy as jnp
2020-03-09 16:04:59 -04:00
Roy Frostig
6da6df0ae1 fix comment typo in MPC/LQR example 2020-02-26 08:00:27 -08:00
Roy Frostig
9cb8171fa7 remove unused imports in MPC/LQR example 2020-02-20 11:07:06 -08:00
Roy Frostig
f02578ea9e set tolerances and toggle dtype checks in MPC/LQR example tests 2020-02-19 19:55:47 -08:00
Roy Frostig
f1af9893a4 move MPC/LQR example to examples directory 2020-02-19 13:47:18 -08: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
Matthew Johnson
8bca2c90e7 fix urllib import for py3 2020-01-09 20:25:42 -08:00
Peter Hawkins
dcc882cf6b
Drop Python 2 support from JAX. (#1962)
Remove six dependency.
2020-01-08 13:17:55 -05:00
Roy Frostig
1648435268 enable kernel regression example test 2019-11-20 14:25:09 -08:00
Jonas Rauber
f8c5d98653 fixed cross-entropy losses in mnist examples (fixes #1023) 2019-10-27 09:43:46 +01:00
Karthik Kumara
1090e89a86 - sign missing from loss function definition 2019-10-17 16:44:27 -07:00
Nicolas Papernot
a411384d48
typo 2019-08-09 17:18:31 -07:00
Fabian Pedregosa
b3df11d64a
Typo in docstring
jax.experimentaloptimizers -> jax.experimental.optimizers
2019-08-06 09:07:48 -04:00
Jasper Snoek
06d41fbad1 Added a note about squared distances 2019-05-30 12:13:10 -04:00
Jasper Snoek
3d1e419e0d Updating GP regression example 2019-05-29 16:36:45 -04:00
Jasper Snoek
58c0221505 Updating GP regression example 2019-05-29 16:35:39 -04:00
Matthew Johnson
3c1b832e81 remove accidentally-checked-in file 2019-05-28 21:55:11 -07:00
Matthew Johnson
7751ba449a remove accidentally-checked-in file 2019-05-28 21:54:13 -07:00
Matthew Johnson
8ffb9417e7 fix stax initialization rng bug, remove temp file 2019-05-23 11:28:15 -07:00
Matthew Johnson
92e5f93a29 tweak docstrings in mnist examples 2019-05-23 09:07:44 -07: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
Elliot Creager
948def817f docstring touchups 2019-04-30 09:26:16 -04:00
Elliot Creager
6d9f465b96 Add Differentially Private SGD example 2019-04-29 16:26:18 -04:00
Matthew Johnson
0cf14837c9 make a lax package, revert control flow names (#607)
c.f. #597
pair=skyewm
2019-04-12 16:28:40 -07:00
Skye Wanderman-Milne
105e46f379 Factor out control flow from lax.py into lax_control_flow.py.
Also moves control flow tests to lax_control_flow_test.py.
2019-04-12 13:57:18 -07:00
Matthew Johnson
4e39876941
Merge pull request #566 from j-towns/jax-random-stax
Use jax.random for stax initialization
2019-04-05 07:45:09 -07:00
Alex Alemi
026a743000
Change MNIST Data URL to CVDF mirror
Change the MNIST data download url to the CVDF mirror, as is done in the main tensorflow and tfds datasets.
2019-04-03 10:18:49 -07:00
Jamie Townsend
70fc02504b Fix examples_test.py 2019-04-03 13:32:14 +01:00
Jamie Townsend
1c9b9a57fd Use jax.random for stax initialization 2019-04-03 12:54:02 +01:00
Matthew Johnson
15da530b03 add spmd mnist example 2019-03-08 09:59:03 -08:00
Peter Hawkins
95483c76e9 Expose logsumexp as scipy.special.logsumexp.
scipy.misc.logsumexp is deprecated and appears slated to be removed entirely in scipy 1.3.
2019-02-24 11:49:15 -05:00
Matthew Johnson
0ff98a74eb add random.fold_in, update mnist_vae.py loops 2019-02-13 09:55:36 -08:00
Roy Frostig
98ace13a98
Merge pull request #355 from JasperSnoek/master
Adding a Gaussian process example
2019-02-11 17:02:16 -08:00
Jasper Snoek
130694affa Adding a Gaussian process example 2019-02-11 16:41:51 -05:00
Matthew Johnson
a26fa98460 fix examples_test 2019-02-11 13:38:04 -08:00
Matthew Johnson
0c2cd5f858 tweaks to make internal tests pass 2019-02-11 13:32:22 -08:00
Matthew Johnson
5847d96168 rename "minmax" -> "optimizers" 2019-02-06 11:02:16 -08:00
He Jia
ed087ddbfa Minor updates of advi.py
1. fixed some syntax errors
2. hold on the plot in the end
2019-01-18 07:38:43 -08:00