106 Commits

Author SHA1 Message Date
Jake VanderPlas
d85d204897 DOC: change build from nbsphinx to myst 2021-02-16 10:28:39 -08:00
Jake VanderPlas
6eedadc27f DOC: use jupytext to maintain synced markdown versions of notebooks. 2021-02-10 12:23:07 -08:00
Jake VanderPlas
7545aeb8d3 doc build: avoid executing thinking_in_jax.ipynb 2021-02-04 09:26:11 -08:00
Jake VanderPlas
667cf9bcb1 Merge problematic cells in thinking_in_jax.ipynb 2021-02-03 09:25:26 -08:00
Jake VanderPlas
113584c8b0 Add hidden sleep statement in attempt to sidestep RTD build issue 2021-02-01 14:44:18 -08:00
Jake VanderPlas
967f3ac7a3 Add thinking_in_jax.ipynb 2021-01-26 12:08:37 -08:00
Peter Hawkins
3ac809ede3 [JAX] Move jax.util to jax._src_util.
PiperOrigin-RevId: 351234602
2021-01-11 14:21:07 -08:00
Qiao Zhang
533b4ac3c4 Updates XLA_in_Python notebook with the latest API. 2021-01-07 10:14:50 -08:00
jax authors
4e0b8a8a10 Merge pull request #4948 from gmittal:master
PiperOrigin-RevId: 346291917
2020-12-08 04:01:02 -08:00
Gautam Mittal
df411ff49b
Update autodiff_cookbook.ipynb 2020-12-07 10:06:08 -08:00
Gautam Mittal
4440ef4568 Move example 2020-12-07 10:05:00 -08:00
jax authors
3a4e35f4e0 Internal change
PiperOrigin-RevId: 346099032
2020-12-07 08:31:40 -08:00
Gautam Mittal
b2a7be9806
Update quickstart.ipynb 2020-11-18 16:22:07 -08:00
Gautam Mittal
f54e531cc4 Add vector-valued gradient example. 2020-11-18 16:18:43 -08:00
Matthew Johnson
491ad9657d fix typo in autodiff cookbook 2020-11-12 14:33:17 -08:00
Matthew Johnson
e51163af32 only pass hashable values as static args 2020-10-16 13:11:56 -07:00
jax authors
4a20eea828 Copybara import of the project:
--
609f6f3e16d21fed34cc5269c54a0d78ac44a8bc by Matthew Johnson <mattjj@google.com>:

fix custom_jvp/vjp closure issues

PiperOrigin-RevId: 337457689
2020-10-16 00:21:32 -07:00
johnpjf
b357005e60
Update Common_Gotchas_in_JAX.ipynb
Clarify that the index is clamped to the bounds of the array when accessing out of bounds.
2020-10-05 10:19:18 -07:00
Qiao Zhang
26a53ae554
Add comments for residuals from f_bwd. (#4244) 2020-09-10 13:58:28 +03:00
George Thomas
b4efb31f1a
Docs: Fix broken link in quickstart (#4102) 2020-08-19 11:36:28 -07:00
Matthew Johnson
09d8ac14de
use fewer internal APIs in custom interp notebook (#4016) 2020-08-10 20:48:03 -07:00
Justin Lebar
e8c7d9e281
s/Three-fry/Threefry/ (#3918)
Per http://www.thesalmons.org/john/random123/
2020-07-30 16:52:44 -07:00
Stephan Hoyer
dd7ab39e4d
Fix formatting in the custom derivatives notebook (#3876)
Sphinx is apparently quite picky about consistent use of headers: you can't
skip a header level. We were getting warnings like "WARNING: Title level
inconsistent" in the docs build, and sub-headers weren't showing up on this
page after the first section.
2020-07-27 22:25:16 -07:00
Jake Vanderplas
05904faf0f
Change onp/np to np/jnp in docs & notebooks (#3760) 2020-07-15 13:17:38 -07:00
igorwilbert
c485a5b04a
Improve jax_debug_nans documentation (#3665) 2020-07-06 09:23:01 +03:00
igorwilbert
3d59a0bcba
Adds information about iterator in fori_loop in The Sharp Bits documentation (#3632) 2020-07-04 15:57:15 +03:00
frederikwilde
42cbe49ce6
Correction a typo of the period of the PRNG. (#3578) 2020-06-26 15:19:57 -07:00
8bitmp3
7f9fc27c48
Another small fix of the link rendering in the Autodiff Cookbook - vmap transformation (#3502) 2020-06-21 13:52:34 -07:00
8bitmp3
f0dff9c19b
Fix a link rendering to Autograd's reverse-mode Jacobian method (#3501) 2020-06-21 11:07:29 -07:00
Matthew Johnson
0c29cc15b9 fix typos 2020-06-14 10:13:56 -07:00
Matthew Johnson
1b88fba57c fix and better explain complex JVPs / VJPs
fixes #3433
2020-06-14 10:09:15 -07:00
Matthew Johnson
d3ccf0a5c7 fix typo in docs 2020-06-09 07:35:07 -07:00
Matthew Johnson
2a6d3f417c fix another docs bug 2020-06-08 14:13:01 -07:00
Matthew Johnson
0a716978ad fix a docs bug 2020-06-08 13:30:00 -07:00
Adam Paszke
7a0e8bae30 Update notebooks 2020-06-05 15:52:03 +00:00
Skye Wanderman-Milne
66ba734882
Add note to docs describing how pytree arguments work. (#3284)
Addresses #3095. I'm not sure if we wanna link to this from API
docstrings.

This also subsumes the original pytrees notebook.
2020-06-03 09:46:00 -07:00
Jake Vanderplas
41292a2af6
mention numpy & scipy convolve functions in gotchas doc. (#3214) 2020-05-27 13:20:27 -07:00
Jascha Sohl-Dickstein
190f88dede
Update Common_Gotchas_in_JAX.ipynb (#3189)
typo fix
2020-05-22 14:12:44 -07:00
Sebastian Bischoff
b6777c0788
Correct calculation of loss and increase learning rate (#3113) 2020-05-19 15:18:34 -07:00
Sandu Ursu
a9c1b38659
Added link to README (#3139) 2020-05-18 14:12:52 -07:00
sracaniere
72cd1f7dfc
Fix sign error in examples. (#3031) 2020-05-11 08:52:55 -07:00
Peter Hawkins
0ea22b7e19
Use a whitelist to restrict visibility in top-level jax namespace. (#2982)
* Use a whitelist to restrict visibility in top-level jax namespace.

The goal of this change is to capture the way the world is (i.e., not break users), and separately we will work on fixing users to avoid accidentally-exported APIs.
2020-05-07 17:24:19 -04:00
yurodiviy
56f6294e37
Implement nanargmin-max and add tests (#2398)
Co-authored-by: vlad <veryfakemail@ya.ru>
2020-04-28 15:23:03 -04:00
Peter Hawkins
5290c03a17
Remove usage of xla_client.{Computation,ComputationBuilder}. (#2808)
* Remove usage of xla_client.{Computation,ComputationBuilder}.

ComputationBuilder is a fairly pointless wrapper class that mimics an outdated version of the the C++ XLA API. It dates back from when we used to have SWIG bindings and needed to write a non-trivial Python shim to keep the interface pleasant to use. Now that we have pybind11-based bindings that are reasonably ergonomic by themselves, we don't need the wrapper class. Instead, we can simply call the pybind11-wrapped C++ API directly, removing the impedance mismatch between the C++ and Python APIs and allowing us to delete the Python ComputationBuilder class.

Similarly we can delete xla_client.Computation for the same reasons; it doesn't do anything useful on top of the C++ API.
2020-04-23 18:30:47 -04:00
Matthew Johnson
964cf4fb6a
autodiff cookbook: assume continuous second derivatives
fixes #2772
2020-04-21 14:08:26 -07:00
Lauro Langosco di Langosco
d6ab70c315
Fix minor typo in cell (#2692)
* Fix minor typo in cell

One of the arguments to `hvp` wasn't being used, which made the example slightly confusing.

* Fix both definitions of hvp in the autodiff cookbook.

Co-authored-by: Peter Hawkins <phawkins@google.com>
2020-04-13 21:26:30 -04:00
Roman Ring
656c3a9504
fix a typo in docs notebook (#2672) 2020-04-10 15:30:01 -04:00
George Necula
abbc70b20a Added type annotations and comments related to partial evaluation.
Introduced two new constructors for PartialVal: unknown and known.
These should make it easier to read the code where we construct
PartialVal:

 * instead of PartialVal((aval, core.unit) we use PartialVal.unknown(aval)
 * instead of PartialVal((None, pval)) we use PartialVal.known(pval)

Also disabled some new tests in random_tests.py on Mac. They segfault,
apparently due to the same issue #432.
2020-04-09 13:00:33 +03:00
Jin Dong
6213f8b81e
Remove unnecessary code in colabs (#2623)
* fix misspell in autodiff_cookbook[modify colab directly]

* remove unnecessary from __future__ code[modify colab directly]

* change tf&tfds-nightly to stable version
2020-04-06 17:26:51 -07:00
Matthew Johnson
a4ceae1c00 fix link in custom derivatives tutorial notebook 2020-03-30 22:12:38 -07:00