398 Commits

Author SHA1 Message Date
Qiao Zhang
26a53ae554
Add comments for residuals from f_bwd. (#4244) 2020-09-10 13:58:28 +03:00
Qiao Zhang
4600dd7957
Update jaxlib version for dlpack fix. (#4231) 2020-09-08 17:20:48 -07:00
Matthew Johnson
7f3078b70d
updtate version and changelog for pypi (#4224) 2020-09-08 08:54:13 -07:00
Alex Riley
708d07d5ff
Add jax.numpy.array_split (#4197) 2020-09-02 16:13:17 -07:00
Matthew Johnson
6b6789a53b
applied simple find+sed for 'master' -> 'main' (#4174)
* applied simple find+sed for 'master' -> 'main'

* Rename master->main in JAX API and internals (#4178)

* Started with #4174 
* Renamed Trace.master to Trace.main
* Renamed core.new_master and core.new_base_master

Co-authored-by: George Necula <gcnecula@gmail.com>
2020-08-30 11:16:51 +03:00
Wojciech Rzadkowski
9d733dd7d7
Doc: change suggested way of starting the profiler (#4120) 2020-08-21 11:19:51 -07:00
George Thomas
b4efb31f1a
Docs: Fix broken link in quickstart (#4102) 2020-08-19 11:36:28 -07:00
gaurav pathak
29f7fa753a
Add implementation of jax.numpy.trim_zeros (#4027) 2020-08-18 13:40:45 -07:00
Roy Frostig
df6a3da449 add scan unrolling to a previous changelog entry 2020-08-13 14:29:18 -07:00
Matthew Johnson
09d8ac14de
use fewer internal APIs in custom interp notebook (#4016) 2020-08-10 20:48:03 -07:00
Jake Vanderplas
d551cec6e8
Add Bessel functions in jax.numpy & jax.scipy.special (#4007) 2020-08-10 10:10:59 -07:00
Jamie Townsend
89e8a0886b
Fix warnings in pmap_test.py (#3977)
Also add note to developer documentation re: testing pmap.
2020-08-10 10:09:34 -07:00
Jake Vanderplas
dd304f62c1
Add jax.numpy.invert (#3998) 2020-08-10 08:41:54 -07:00
Justin Lebar
fb02cc80f7
Update docs requirements. (#3969) 2020-08-07 11:47:05 -07:00
Justin Lebar
a98103f678
Add missing functions to jax.numpy docs (#3981) 2020-08-06 12:57:19 -07:00
Stephan Hoyer
816bcd7196
Fix jax.checkpoint in API docs (#3980)
On the index API doc page, it turns out functions need to be listed *twice* to
appear.
2020-08-06 09:39:33 -07:00
Peter Hawkins
56a55b1749
Fix formatting problem in GPU memory allocation docs. (#3979) 2020-08-06 12:00:56 -04:00
Justin Lebar
2277728e8c
Add numpy.vectorize to list of numpy fns in docs (#3968) 2020-08-05 18:37:29 -07:00
Jake Vanderplas
d6f131aaee
implement jax.numpy.interp (#3949) 2020-08-04 12:39:04 -07:00
Jake Vanderplas
e0ebb144f9
Add switch and associative_scan to lax docs (#3946) 2020-08-03 12:32:32 -07:00
Jake Vanderplas
1229f44adf
add jnp.piecewise implementation (#3943) 2020-08-03 12:31:52 -07:00
Stephen Tu
3b68fab2d8
Update faq.rst (#3934) 2020-08-01 11:12:59 -07:00
Matthew Johnson
abfd33570b
update version and changelog for pypi (#3921)
* update version and changelog for pypi

* fix typos
2020-07-30 22:19:29 -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
Matthew Johnson
de645c5b8b
update version and changelog for pypi (#3906) 2020-07-29 22:32:49 -07:00
Jake Vanderplas
e0a8d44861
Add jnp.modf() & improve test coverage for related functions (#3894) 2020-07-29 12:53:28 -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
Vaibhav Srivastav
3aa37d3af4
Replicating sort_complex functionality from np.sort_complex to jax.numpy (#3870) 2020-07-27 14:27:36 -07:00
Jake Vanderplas
17c21a79c3
update jaxlib version and changelog for pypi (#3824) 2020-07-22 15:21:30 -07:00
Jake Vanderplas
edff6933f8
update jax version and changelog for pypi (#3823) 2020-07-22 14:20:34 -07:00
Jake Vanderplas
04a6238d4e
implement jax.numpy.lexsort (#3812) 2020-07-22 12:48:49 -07:00
Jake Vanderplas
05904faf0f
Change onp/np to np/jnp in docs & notebooks (#3760) 2020-07-15 13:17:38 -07:00
Roy Frostig
8a62a9b654
block-unrolled scan primitive implementation (#3738)
* block-unrolled scan implementation, via optional `_unroll` scan parameter

* index statically in the inlined path of lax.scan

* make `unroll` a required scan parameter, and test that it unrolls
2020-07-15 14:00:50 -04:00
Chase Roberts
2b7a39f92b
Add pshuffle to docs (#3742) 2020-07-14 09:05:45 -04:00
Alex Dragan
0d81e988d8
Implement np.intersect1d (#3726)
* Implement np.intersect1d

* Add jitable helper to function

* Fix argsort failing tests

* Fix linter errors
2020-07-12 22:32:41 -07:00
Alex Dragan
412b9d5209
hfft and ihfft implementation (#3664) 2020-07-10 10:34:59 -07:00
Jake Vanderplas
d2f9c46a0c
Remove some non-inclusive language (#3710) 2020-07-10 09:29:06 -07:00
Peter Hawkins
b943b31b22
Add jax.image.resize. (#3703)
* Add jax.image.resize.

This is a port of `tf.image.resize()` and the `ScaleAndTranslate` operator.

While I don't expect this implementation to be particularly fast, it is a useful generic implementation to which we can add optimized special cases as the need arises.
2020-07-10 09:57:59 -04:00
8bitmp3
11b40fb4b2
Fix a link to TensorBoard's profiler in Profiling JAX Programs (#3692) 2020-07-08 10:53:06 -07:00
8bitmp3
242b382bab
Remove a deprecated reference to testExamplesJaxprDoc in Understanding Jaxpr (#3680) 2020-07-07 11:29:44 -07:00
Peter Hawkins
a2fb798dc3
Rename swish to silu, add swish as an alias to silu. (#3673) 2020-07-06 18:08:16 -04:00
igorwilbert
c485a5b04a
Improve jax_debug_nans documentation (#3665) 2020-07-06 09:23:01 +03:00
8bitmp3
3bbe06e81a
Add an empty line to JAX change log (#3661) 2020-07-05 11:26:50 +03:00
George Necula
4f3011f320
Refactored host_callback to use the C++ runtime. (#3644)
* Refactored host_callback to use the C++ runtime.

* The new runtime makes it unnecessary to start the outfeed_receiver
  in the user's code
* We don't need msgpack anymore
* There is an interaction between host_callback and using lax.outfeed.
  I am trying to solve this by (a) making host_callback_test stop the
  outfeed receiver on finish and infeed_test on start, and (b)
  telling pytest-xdist to run all the tests from one file into
  a single worker.
2020-07-04 18:12:58 +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
George Necula
a5ed161550
Release jaxlib 0.1.51 (#3636) 2020-07-02 10:33:00 +03:00
Peter Hawkins
141fabbbf5
Reimplement argmin/argmax using a single pass variadic reduction. (#3611) 2020-07-01 11:01:22 -04:00
Jake VanderPlas
afce718eb1 Add ability to specify individual test targets 2020-06-29 11:08:57 -07:00
Matthew Johnson
99cfbfb7c3
bump version for pypi (#3590) 2020-06-28 14:47:39 -07:00
Sri Hari Krishna Narayanan
7b57dc8c80
Issue1635 expm frechet (#2062)
* Implement Frechet derivatives for expm.

* Update expm to use the current custom gradients API.

Make some stylistic fixes.

Co-authored-by: Peter Hawkins <phawkins@google.com>
2020-06-28 12:11:12 -04:00