15 Commits

Author SHA1 Message Date
Roy Frostig
78fd4f1664 update top-level examples to use new-style typed keys 2024-03-07 12:40:10 -08: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
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
Roy Frostig
9f8e3df320 fix advi example: pass tuples rather than lists to jitted function 2021-08-10 10:31:26 -07:00
Peter Hawkins
26e9ebcdae Move jax.api to jax._src.api.
PiperOrigin-RevId: 368233837
2021-04-13 09:43:24 -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
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
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
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
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
Matthew Johnson
75958898f9 add copyright header to advi.py 2019-01-07 08:59:56 -08:00
David Duvenaud
cb7d3550ff Updated advi example to map over rngs. 2019-01-06 09:16:29 -04:00
David Duvenaud
d85c0dbecb Added example of automatic differentiation variational inference 2018-12-31 16:08:23 -05:00