1689 Commits

Author SHA1 Message Date
George Necula
c4e4199508 Fix typo in notebook JSON (#1441)
Perhaps this is the reason why readthedocs does not show the link to the notebook?
2019-10-04 08:38:36 -07:00
George Necula
73144fafb5
Update docs/notebooks/How_JAX_primitives_work.ipynb
Co-Authored-By: James Bradbury <jekbradbury@google.com>
2019-10-04 11:58:07 +02:00
George Necula
ef886c7fcc
Update docs/notebooks/How_JAX_primitives_work.ipynb
Co-Authored-By: James Bradbury <jekbradbury@google.com>
2019-10-04 11:56:06 +02:00
Roy Frostig
db3b0dd036 fix missing delimiter in notebook source 2019-10-03 14:02:45 -07:00
George Necula
0fdf377343 Switched testing to using sphinx (part of the documentation generation) 2019-10-03 20:54:28 +02:00
George Necula
57e39ce938 Added support for testing Colab notebooks.
Testing is done by running "jupyter nbconvert --to notebook" and
then parsing the resulting notebook to look for errors.

One can declare expected errors, and the test will fail if those
are missing.

In the process of doig this, found and fixed a bug in the autodiff_cookbook
notebook.
2019-10-03 20:07:49 +02:00
George Necula
7bc0f70f87 Fix the title to show nicely in sphinx 2019-10-03 11:59:51 +02:00
George Necula
75c2236063 Addressed comments for the Colab.
* Cleaned up use of section levels
* Renamed ma to multiply_add and sq_add to square_add
* Other minor clarifications
* Separated the Colabs into Tutorials and Advanced Tutorials
2019-10-03 11:20:04 +02:00
George Necula
454320e9c9 Added How_JAX_primitives_work colab 2019-10-02 14:42:01 +02:00
Stephan Hoyer
08254fe5bc Fixup colab notebook URLs 2019-09-30 11:12:01 -07:00
Stephan Hoyer
9bd7330e1f
Notebooks on RTD (#1121) 2019-09-30 11:00:02 -07:00
Peter Hawkins
454cb4f6e9 Implement jax.scipy.linalg.lu_solve. 2019-09-05 09:59:47 -04:00
James Bradbury
146b5d121a
Merge pull request #1262 from google/jb/initializers
Migrate initializers and activation functions to jax.nn
2019-09-04 14:48:20 -07:00
Matthew Johnson
c5737d0f6c
Merge pull request #1272 from alexalemi/vectorize
Convert gufunc notebook to experimental module.
2019-09-02 07:26:41 -07:00
Matthew Johnson
cac042c34a move asinh/acosh/atanh to lax_numpy.py only 2019-08-31 22:39:51 -07:00
Alex Alemi
55a06cbd5b added documentation 2019-08-30 21:23:18 +00:00
James Bradbury
cc49d8b325 add docs for jax.nn 2019-08-29 18:15:36 -07:00
Matthew Johnson
c6f430a831 tweak rank promotion warning docs 2019-08-25 14:36:23 -07:00
Matthew Johnson
94fafe880c attempt to fix sphinx rst rendering 2019-08-25 14:33:14 -07:00
Matthew Johnson
3c2a73592c improve rank promotion warning, add doc page 2019-08-25 14:28:53 -07:00
Matthew Johnson
1e375dafe0
Merge pull request #1217 from shoyer/doc-index-page
Fix index of API docs
2019-08-21 16:46:29 -07:00
Skye Wanderman-Milne
c50918feed Update GPU memory documentation.
See #1222.
2019-08-21 12:02:42 -07:00
Stephan Hoyer
ee54f6553c Fix index of API docs
The current index of the API docs page seems to have broken links: when I
click on "Automatic differentiation" for example, I get sent to the "JIT"
section.

This change fixes the links.
2019-08-20 23:13:15 -07:00
Stephan Hoyer
c3bf532aae Fix sphinx references in JAX docs
The links in `concurrency.rst` should work now.

The links in async_dispatch.rst don't work yet, but they will if
``DeviceArray`` ever ends up JAX's API docs.
2019-08-10 16:40:08 -07:00
Peter Hawkins
8fd6e19ef6 More concurrency fixes. 2019-08-09 15:08:26 -04:00
Peter Hawkins
6dc730a5f4 Make JAX tracer state thread-local. Allows performing traces in separate threads.
Using threading within a traced context still won't work, but that is perhaps less important than the ability to call JIT-ted computations from separate threads.

(Revives https://github.com/google/jax/pull/734.)
2019-08-09 13:55:20 -04:00
Peter Hawkins
bcacdfe315 Add some brief documentation about how to profile/trace JAX programs. 2019-08-08 21:02:41 -04:00
Peter Hawkins
6337ac611a Mention tf.config.experimental.set_visible_devices in GPU memory allocation docs. 2019-08-06 16:51:40 -04:00
Matthew Johnson
6e5dae4bd4
Merge pull request #1119 from shoyer/docs-jax-logo
Add the JAX logo to sphinx docs
2019-08-05 13:47:24 -07:00
Stephan Hoyer
515e07bdf2 White background 2019-08-05 13:08:17 -07:00
Stephan Hoyer
8c8222a319 Add the JAX logo to sphinx docs 2019-08-05 12:34:29 -07:00
Stephan Hoyer
8c628a267b Implement lax.map
Fixes GH-1113
2019-08-05 12:14:05 -07:00
Peter Hawkins
b45d1ec1dd Fix test breakage at head.
Add new numpy/scipy functions to documentation.
2019-08-02 08:55:22 -04:00
Peter Hawkins
6434d9d029 Add corrcoef to jax.numpy docs. 2019-07-30 21:23:44 -04:00
Skye Wanderman-Milne
3461dfffab Add how to configure TF to GPU memory allocation doc. 2019-07-30 15:52:05 -07:00
Skye Wanderman-Milne
5b9849f177 Add docs on GPU memory allocation. 2019-07-30 15:49:35 -07:00
Peter Hawkins
4eb1820ae2 Add documentation to JAX modules. 2019-07-21 15:55:47 -04:00
Peter Hawkins
586385bce0 Add cov, percentile, and quantile to the jax.numpy docs. 2019-07-08 13:37:32 -04:00
Matthew Johnson
0f1d913db3 add xla_computation to jax rst docs 2019-07-05 16:55:11 -07:00
Peter Hawkins
59be9b7a04 Minor doc fixes. 2019-07-02 15:00:47 -04:00
Peter Hawkins
2ef945db64 Add numpy.select to the documentation. 2019-06-24 10:51:21 -04:00
Justin Lebar
3da993802b Update docs to include index_{min,max}. 2019-06-23 15:32:43 -07:00
Peter Hawkins
0190684ee2
Merge pull request #866 from hawkinsp/master
Implement np.ix_, for non-bool inputs.
2019-06-17 18:39:45 -06:00
Peter Hawkins
2dc89ad4a6 Add ix_ to the docs. 2019-06-17 17:19:21 -04:00
Matthew Johnson
bbf625e0df fix jax.rst docs (remove defvjp2 / defjvp2) 2019-06-11 06:52:55 -07:00
Matthew Johnson
ab20f0292c add docstring for defjvp_all 2019-06-05 17:34:14 -07:00
Matthew Johnson
720dec4072 add custom_gradient 2019-06-05 13:48:04 -07:00
Peter Hawkins
91e6ba322c Add documentation on asynchronous dispatch. 2019-06-04 10:09:43 -04:00
Matthew Johnson
ffec059f0e add jax.eval_shape to reference docs via jax.rst 2019-06-01 09:53:32 -07:00
Matthew Johnson
2f645dd36d remove one lax parallel op from sphinx docs 2019-05-17 13:24:06 -07:00