| 25b4070 | jax authors | 26 October 2020, 22:17:45 UTC | Merge pull request #4707 from google:traceback-register PiperOrigin-RevId: 339129121 | 26 October 2020, 22:17:45 UTC |
| 4c9545d | jax authors | 26 October 2020, 21:01:37 UTC | Merge pull request #4709 from google:update-pypi PiperOrigin-RevId: 339112147 | 26 October 2020, 21:01:37 UTC |
| 3bb0c49 | jax authors | 26 October 2020, 20:58:11 UTC | Merge pull request #4702 from fabianp:patch-3 PiperOrigin-RevId: 339111474 | 26 October 2020, 20:58:11 UTC |
| 4ddffff | Matthew Johnson | 26 October 2020, 20:20:32 UTC | update version and changelog for pypi | 26 October 2020, 20:20:32 UTC |
| fa1d7ab | Roy Frostig | 26 October 2020, 19:30:51 UTC | move `wraps` from `api_util` to `util` to avoid cyclic dependencies | 26 October 2020, 19:31:19 UTC |
| 5d50e19 | Roy Frostig | 26 October 2020, 17:03:06 UTC | add path exclusion opt-in to filtered stack traces and use it throughout the codebase | 26 October 2020, 19:31:19 UTC |
| abbd827 | jax authors | 24 October 2020, 18:42:17 UTC | Merge pull request #4698 from google:custom-gradient-wrapper PiperOrigin-RevId: 338846354 | 24 October 2020, 18:42:17 UTC |
| 5d73628 | Fabian Pedregosa | 24 October 2020, 14:00:23 UTC | FIX: example usage of optimizers The example usage of the optimizers API in this docstring did not correspond with the current API where optimizers return a triplet (opt_init, opt_update, get_params) and not an object with .init, .update and .get_params fields. Using the previous code would result in the error "AttributeError: 'Optimizer' object has no attribute 'init'" | 24 October 2020, 14:00:23 UTC |
| 3c6cdcf | Matthew Johnson | 23 October 2020, 20:54:23 UTC | add jax.custom_gradient wrapper for jax.custom_vjp There was a deprecatd version of this wrapper implemented in terms of jax.custom_transforms (which itself is deprecated, and hopefully soon to be removed), but this commit adds an implementation in terms of jax.custom_vjp. One drawback it has relative to jax.custom_vjp is that it doesn't support Python control flow in the backward-pass function. | 24 October 2020, 05:32:19 UTC |
| 602c783 | jax authors | 23 October 2020, 23:52:32 UTC | Merge pull request #4700 from hawkinsp:fft PiperOrigin-RevId: 338769487 | 23 October 2020, 23:52:32 UTC |
| 9b7fd63 | Peter Hawkins | 23 October 2020, 23:38:43 UTC | Add PocketFFT mirror to fix download problems for some Bazel versions. | 23 October 2020, 23:48:00 UTC |
| f58f1ee | Peter Hawkins | 23 October 2020, 21:20:06 UTC | [JAX] Use PocketFFT for FFTs on CPU instead of Eigen. PocketFFT is the same FFT library used by NumPy (although we are using the C++ variant rather than the C variant.) For the benchmark in #2952 on my workstation: Before: ``` 907.3490574884647 max: 4.362646594533903e-08 mean: 6.237288307614869e-09 min: 0.0 numpy fft execution time [ms]: 37.088446617126465 jax fft execution time [ms]: 74.93342399597168 ``` After: ``` 907.3490574884647 max: 1.9057386696477137e-12 mean: 3.9326737908882566e-13 min: 0.0 numpy fft execution time [ms]: 37.756404876708984 jax fft execution time [ms]: 28.128278255462646 ``` Fixes https://github.com/google/jax/issues/2952 PiperOrigin-RevId: 338743753 | 23 October 2020, 21:20:32 UTC |
| 8121255 | jax authors | 23 October 2020, 21:12:13 UTC | Merge pull request #4661 from jblespiau:changelist/338124676 PiperOrigin-RevId: 338742268 | 23 October 2020, 21:12:13 UTC |
| 98ee69b | Jean-Baptiste Lespiau | 20 October 2020, 20:22:33 UTC | Emit a warning on non-hashble static arguments in jax.jit. The message looks like, e.g.: Static argument (index 1) of type <class 'numpy.ndarray'> for function f is non-hashable. As this can lead to unexpected cache-misses, it will raise an error in a near future. | 23 October 2020, 20:29:41 UTC |
| 9ba28d2 | Matthew Johnson | 23 October 2020, 14:34:32 UTC | Copybara import of the project: -- ced333d1d4aec2825e9afd81c2ca9721b7e3cc67 by Matthew Johnson <mattjj@google.com>: redo #4535 lazy simplification PiperOrigin-RevId: 338670328 | 23 October 2020, 14:35:01 UTC |
| fcaced3 | Matthew Johnson | 23 October 2020, 04:17:56 UTC | Copybara import of the project: -- ced333d1d4aec2825e9afd81c2ca9721b7e3cc67 by Matthew Johnson <mattjj@google.com>: redo #4535 lazy simplification PiperOrigin-RevId: 338606348 | 23 October 2020, 04:18:22 UTC |
| 2677c04 | jax authors | 23 October 2020, 01:16:02 UTC | Merge pull request #4658 from jakevdp:custom-obj-avals PiperOrigin-RevId: 338588156 | 23 October 2020, 01:16:02 UTC |
| 8ff6396 | jax authors | 23 October 2020, 00:29:55 UTC | Merge pull request #4686 from google:lazy-simplification-again PiperOrigin-RevId: 338582074 | 23 October 2020, 00:29:55 UTC |
| fe91728 | jax authors | 23 October 2020, 00:02:26 UTC | Merge pull request #4685 from google:make-dynamic-slice-transpose-handle-symbolic-zero-cotangents PiperOrigin-RevId: 338577125 | 23 October 2020, 00:02:26 UTC |
| 9845cce | jax authors | 22 October 2020, 23:59:04 UTC | Merge pull request #4684 from jakevdp:cleanup PiperOrigin-RevId: 338577036 | 22 October 2020, 23:59:04 UTC |
| ced333d | Matthew Johnson | 22 October 2020, 23:37:08 UTC | redo #4535 lazy simplification | 22 October 2020, 23:56:29 UTC |
| a8b9056 | jax authors | 22 October 2020, 22:33:03 UTC | Merge pull request #4681 from jblespiau:changelist/338538921 PiperOrigin-RevId: 338562006 | 22 October 2020, 22:33:03 UTC |
| f40ac06 | Matthew Johnson | 22 October 2020, 22:31:43 UTC | make lax.dynamic_slice transpose handle symb zeros | 22 October 2020, 22:31:43 UTC |
| 2af475d | Jake VanderPlas | 22 October 2020, 22:18:03 UTC | Cleanup: pass function name rather than function object | 22 October 2020, 22:18:03 UTC |
| 225cdb9 | Jean-Baptiste Lespiau | 22 October 2020, 20:34:04 UTC | Fix a bug in `make_jaxpr`, as static-arguments should be passed as is | 22 October 2020, 20:35:36 UTC |
| 48db25e | Jake VanderPlas | 20 October 2020, 19:42:17 UTC | [multi-buf] simplify custom object test avals | 22 October 2020, 16:42:38 UTC |
| 20ee415 | jax authors | 22 October 2020, 16:12:32 UTC | Merge pull request #4644 from jblespiau:changelist/337949790 PiperOrigin-RevId: 338484804 | 22 October 2020, 16:12:32 UTC |
| 62cbc23 | Jean-Baptiste Lespiau | 22 October 2020, 15:57:12 UTC | Raise a meaningful exception in the C++ jax.jit on non-hashable object static argnum. This also fixes 2 tests that were passing non-hashable values to a jitted function. PiperOrigin-RevId: 338481880 | 22 October 2020, 15:57:46 UTC |
| cbd0779 | jax authors | 21 October 2020, 23:03:13 UTC | Merge pull request #4635 from aldragan0:numpy-setdiff1d PiperOrigin-RevId: 338358849 | 21 October 2020, 23:03:13 UTC |
| b7fee45 | Alex Dragan | 17 October 2020, 20:48:39 UTC | [impl] Add support for setdiff1d | 21 October 2020, 17:17:12 UTC |
| 72a273f | jax authors | 21 October 2020, 16:44:25 UTC | Merge pull request #4459 from romanngg:extract_patches PiperOrigin-RevId: 338281178 | 21 October 2020, 16:44:25 UTC |
| bcc1d71 | jax authors | 21 October 2020, 08:53:46 UTC | Merge pull request #4650 from gnecula:abstract_check PiperOrigin-RevId: 338222103 | 21 October 2020, 08:53:46 UTC |
| da0bff2 | Roman Novak | 21 October 2020, 05:58:53 UTC | Add `lax.conv_general_dilated_patches` | 21 October 2020, 05:58:53 UTC |
| c5b983c | George Necula | 21 October 2020, 04:53:37 UTC | Update jax/core.py Co-authored-by: Roy Frostig <froystig@users.noreply.github.com> | 21 October 2020, 04:53:37 UTC |
| e098fbd | jax authors | 21 October 2020, 04:42:51 UTC | Merge pull request #4665 from google:fix-custom-jvp-vmap-bug PiperOrigin-RevId: 338197181 | 21 October 2020, 04:42:51 UTC |
| 32ffb24 | Matthew Johnson | 21 October 2020, 04:20:04 UTC | fix flake | 21 October 2020, 04:20:04 UTC |
| d5c9408 | Matthew Johnson | 21 October 2020, 04:16:00 UTC | improve tests | 21 October 2020, 04:16:00 UTC |
| a46d002 | Matthew Johnson | 21 October 2020, 04:08:59 UTC | fix a custom_jvp vmap bug from @dpfau | 21 October 2020, 04:08:59 UTC |
| dcfa764 | jax authors | 21 October 2020, 02:22:02 UTC | Merge pull request #4663 from google:improve-tracer-error-2 PiperOrigin-RevId: 338183176 | 21 October 2020, 02:22:02 UTC |
| 2b62a44 | Matthew Johnson | 21 October 2020, 02:10:41 UTC | only run test with omnistaging on | 21 October 2020, 02:10:41 UTC |
| 49b1ec2 | Matthew Johnson | 21 October 2020, 01:26:32 UTC | add test | 21 October 2020, 01:26:32 UTC |
| 8684da4 | jax authors | 21 October 2020, 01:18:11 UTC | Merge pull request #4664 from google:4008-fixes PiperOrigin-RevId: 338176122 | 21 October 2020, 01:18:11 UTC |
| 79e3db5 | Matthew Johnson | 21 October 2020, 00:51:51 UTC | fixes on #4008 (thanks @apaszke) | 21 October 2020, 00:51:51 UTC |
| 0090502 | jax authors | 21 October 2020, 00:49:35 UTC | Merge pull request #4662 from google:skip-dlpack-tests-on-jaxlib-version PiperOrigin-RevId: 338172795 | 21 October 2020, 00:49:35 UTC |
| ac6cd23 | Matthew Johnson | 20 October 2020, 23:10:56 UTC | improve an escaped tracer error | 20 October 2020, 23:10:56 UTC |
| 4a19de6 | Matthew Johnson | 20 October 2020, 22:35:03 UTC | added test skips based on jaxlib version (o/w these were failing at HEAD for me) | 20 October 2020, 22:35:03 UTC |
| d001ac6 | Peter Hawkins | 20 October 2020, 20:06:37 UTC | [JAX] Add support for retaining ownership of DLPack tensors. Move dlpack.py contents under jax/_src/dlpack.py. Add array interoperability test between JAX and TensorFlow using DLPack. Fixes: https://github.com/google/jax/issues/4636 PiperOrigin-RevId: 338120910 | 20 October 2020, 20:07:07 UTC |
| 41eea2c | jax authors | 20 October 2020, 19:40:54 UTC | Merge pull request #4657 from skye:jaxlib2 PiperOrigin-RevId: 338114862 | 20 October 2020, 19:40:54 UTC |
| cb591eb | George Necula | 20 October 2020, 19:28:22 UTC | Update jax/core.py Co-authored-by: Roy Frostig <froystig@users.noreply.github.com> | 20 October 2020, 19:28:22 UTC |
| a082169 | Skye Wanderman-Milne | 14 October 2020, 17:30:47 UTC | Update README and jaxlib.__version__ for new jaxlib release | 20 October 2020, 18:46:48 UTC |
| 52dbeee | jax authors | 20 October 2020, 18:17:48 UTC | Merge pull request #4656 from hawkinsp:xla PiperOrigin-RevId: 338096915 | 20 October 2020, 18:17:48 UTC |
| 29a83dd | Peter Hawkins | 20 October 2020, 18:11:47 UTC | Update XLA. Incorporates fix for #4651. Fixes #4651. | 20 October 2020, 18:11:47 UTC |
| 9c0452d | jax authors | 20 October 2020, 17:55:30 UTC | Merge pull request #4655 from skye:cpp_jit PiperOrigin-RevId: 338092016 | 20 October 2020, 17:55:30 UTC |
| 8ef49a3 | Skye Wanderman-Milne | 20 October 2020, 16:52:00 UTC | Don't turn on C++ jit by default until we're sure it works everywhere. | 20 October 2020, 16:52:00 UTC |
| 537427a | Alexey Radul | 20 October 2020, 15:24:53 UTC | Copybara import of the project: -- 138105a9ea44e7a8c3ce575a4e51b7ed51518d41 by Skye Wanderman-Milne <skyewm@google.com>: Update README, CHANGELOG, and jaxlib.__version__ for new jaxlib release PiperOrigin-RevId: 338063494 | 20 October 2020, 15:25:23 UTC |
| 455f6bd | jax authors | 20 October 2020, 14:10:01 UTC | Merge pull request #4652 from hawkinsp:stats PiperOrigin-RevId: 338053891 | 20 October 2020, 14:10:01 UTC |
| c0b480b | Peter Hawkins | 20 October 2020, 13:47:35 UTC | Add missing jax.scipy.stats distributions to the docs. Alphabetize import order. | 20 October 2020, 13:47:35 UTC |
| 114c1dd | Jean-Baptiste Lespiau | 19 October 2020, 22:43:04 UTC | Also convert numpy arrays to hashable values. This was working for a combination of lists and tuples, but the user can also pass a numpy array. Iterable is the correct type here, not Sequence. | 20 October 2020, 11:57:27 UTC |
| 09653bf | George Necula | 20 October 2020, 08:08:39 UTC | Ensure that check_jaxpr is done with abstract values Prior to this it was possible, e.g., for code that contains a Literal, such as to result in FLOPS during checking. The assertion is broken by many tests unless we raise_to_shape for Literals. I have timed the checks on my laptop and I do not see a reduction in the total test time. | 20 October 2020, 08:08:41 UTC |
| 4b8334a | jax authors | 20 October 2020, 02:13:29 UTC | Merge pull request #4647 from google:update-pypi PiperOrigin-RevId: 337980411 | 20 October 2020, 02:13:29 UTC |
| 87623ea | Matthew Johnson | 20 October 2020, 01:55:09 UTC | update version and changelog for pypi | 20 October 2020, 01:55:09 UTC |
| e567b22 | jax authors | 20 October 2020, 00:39:45 UTC | Merge pull request #4571 from skye:warning PiperOrigin-RevId: 337970302 | 20 October 2020, 00:39:45 UTC |
| 60141d9 | jax authors | 19 October 2020, 22:45:56 UTC | Merge pull request #4580 from skye:jaxlib PiperOrigin-RevId: 337950285 | 19 October 2020, 22:45:56 UTC |
| 6a54f16 | jax authors | 19 October 2020, 22:36:35 UTC | Merge pull request #4641 from jakevdp:indexing-deprecation PiperOrigin-RevId: 337948469 | 19 October 2020, 22:36:35 UTC |
| c17fbe2 | jax authors | 19 October 2020, 22:18:17 UTC | Merge pull request #4642 from lmmx:master PiperOrigin-RevId: 337944772 | 19 October 2020, 22:18:17 UTC |
| c56f76d | Jake VanderPlas | 19 October 2020, 21:07:40 UTC | Add deprecation warning for indexing with non-tuple sequences | 19 October 2020, 21:07:40 UTC |
| 11e6f49 | Louis Maddox | 19 October 2020, 20:35:18 UTC | A couple of typo/gap fixes in PRNG design notes | 19 October 2020, 20:35:18 UTC |
| b7ec636 | jax authors | 19 October 2020, 15:27:25 UTC | Merge pull request #4616 from bchetioui:test_linear_solve PiperOrigin-RevId: 337859240 | 19 October 2020, 15:27:25 UTC |
| a8cfc74 | Benjamin Chetioui | 16 October 2020, 10:51:47 UTC | [jax2tf] Implemented a conversion for linear_solve_p | 19 October 2020, 15:07:45 UTC |
| f785794 | jax authors | 19 October 2020, 12:27:48 UTC | Merge pull request #4613 from bchetioui:test_lu PiperOrigin-RevId: 337836222 | 19 October 2020, 12:27:48 UTC |
| 7d30187 | Benjamin Chetioui | 16 October 2020, 13:22:37 UTC | [jax2tf] Implemented conversion of lu_p using _convert_jax_impl. | 19 October 2020, 11:47:41 UTC |
| 26771b3 | jax authors | 17 October 2020, 20:53:48 UTC | Merge pull request #4634 from hawkinsp:src2 PiperOrigin-RevId: 337681899 | 17 October 2020, 20:53:48 UTC |
| e5e33f5 | jax authors | 17 October 2020, 20:50:44 UTC | Merge pull request #4633 from hawkinsp:src PiperOrigin-RevId: 337681883 | 17 October 2020, 20:50:44 UTC |
| 10b7d7d | Peter Hawkins | 17 October 2020, 18:33:26 UTC | Move implementation of jax.lax into jax._src.lax. Remove lax_ prefixes from jax/_src/lax filenames, since they aren't needed any longer to avoid name conflicts. | 17 October 2020, 20:09:21 UTC |
| b078483 | Peter Hawkins | 17 October 2020, 18:51:39 UTC | Import jax.nn.functions by default to fix breakage. | 17 October 2020, 18:51:39 UTC |
| 3ddd390 | Peter Hawkins | 17 October 2020, 18:42:40 UTC | Move jax.third_party to jax._src.third_party. PiperOrigin-RevId: 337675377 | 17 October 2020, 18:43:07 UTC |
| 09dfd47 | jax authors | 17 October 2020, 18:05:40 UTC | Merge pull request #4631 from hawkinsp:src PiperOrigin-RevId: 337673701 | 17 October 2020, 18:05:40 UTC |
| b2808dc | Peter Hawkins | 17 October 2020, 15:31:19 UTC | Move jax.nn implementation into jax._src.nn. | 17 October 2020, 17:45:01 UTC |
| e9909ce | jax authors | 17 October 2020, 17:36:09 UTC | Copybara import of the project: -- a396cfbbd414f6f21f0c7e8a68e6e89d202c0e84 by Peter Hawkins <phawkins@google.com>: Move jax.nn implementation into jax._src.nn. PiperOrigin-RevId: 337671917 | 17 October 2020, 17:40:21 UTC |
| 15bd7f6 | jax authors | 17 October 2020, 17:31:16 UTC | Merge pull request #4628 from hawkinsp:src3 PiperOrigin-RevId: 337671365 | 17 October 2020, 17:31:16 UTC |
| 8841dd4 | jax authors | 17 October 2020, 17:30:56 UTC | Merge pull request #4627 from hawkinsp:src2 PiperOrigin-RevId: 337671357 | 17 October 2020, 17:30:56 UTC |
| ce51b97 | jax authors | 17 October 2020, 17:27:42 UTC | Merge pull request #4626 from hawkinsp:src PiperOrigin-RevId: 337671310 | 17 October 2020, 17:27:42 UTC |
| b42dd89 | jax authors | 17 October 2020, 17:27:25 UTC | Merge pull request #4624 from romanngg:patch-1 PiperOrigin-RevId: 337671286 | 17 October 2020, 17:27:25 UTC |
| e4e54f7 | jax authors | 17 October 2020, 17:23:41 UTC | Merge pull request #4625 from jblespiau:changelist/337664495 PiperOrigin-RevId: 337671242 | 17 October 2020, 17:23:41 UTC |
| 1c7d726 | Jean-Baptiste Lespiau | 17 October 2020, 15:19:11 UTC | Make pad more robust to non-hashable values. The `pad_width` can be an iterable of iterbles, thus, we need to cast both the inner and outer list to tuples. Also, the values should be intergers, but some use-cases (mostly Tensorflow) pass e.g. DeviceArray(1) as values in there. | 17 October 2020, 16:22:45 UTC |
| ef57858 | Peter Hawkins | 17 October 2020, 15:45:28 UTC | Move jax.ops implementation into jax._src.ops. | 17 October 2020, 15:45:28 UTC |
| c876efc | Peter Hawkins | 17 October 2020, 15:37:55 UTC | Move implementation of jax.image to jax._src.image. | 17 October 2020, 15:37:55 UTC |
| a396cfb | Peter Hawkins | 17 October 2020, 15:31:19 UTC | Move jax.nn implementation into jax._src.nn. | 17 October 2020, 15:31:19 UTC |
| 2e6c653 | jax authors | 17 October 2020, 04:36:55 UTC | Merge pull request #4623 from google:issue4622 PiperOrigin-RevId: 337626097 | 17 October 2020, 04:36:55 UTC |
| 9cf018d | Matthew Johnson | 17 October 2020, 04:00:18 UTC | add tests for issue #4622 | 17 October 2020, 04:00:18 UTC |
| 52282cd | Matthew Johnson | 17 October 2020, 03:48:57 UTC | fix test | 17 October 2020, 03:48:57 UTC |
| 621b372 | Roman Novak | 17 October 2020, 03:09:20 UTC | Make `ConvGeneralDilatedDimensionNumbers` type public | 17 October 2020, 03:09:20 UTC |
| ab21409 | jax authors | 17 October 2020, 02:16:22 UTC | Merge pull request #4621 from hawkinsp:src PiperOrigin-RevId: 337616739 | 17 October 2020, 02:16:22 UTC |
| a493a0f | Matthew Johnson | 17 October 2020, 01:21:01 UTC | ensure ConcreteArray equality stays in Python | 17 October 2020, 01:21:01 UTC |
| 1ab0c3d | Matthew Johnson | 17 October 2020, 00:57:59 UTC | refine linearize aval checks | 17 October 2020, 00:57:59 UTC |
| aa107cf | Peter Hawkins | 16 October 2020, 22:08:20 UTC | Move jax.numpy internals into jax._src.numpy. | 17 October 2020, 00:35:19 UTC |
| 9ea1311 | jax authors | 16 October 2020, 22:37:57 UTC | Merge pull request #4620 from sharadmv:callback-custom-jvp PiperOrigin-RevId: 337590572 | 16 October 2020, 22:37:57 UTC |
| 6038b6d | jax authors | 16 October 2020, 22:27:40 UTC | Merge pull request #3712 from google:hashable-static-args-only PiperOrigin-RevId: 337588918 | 16 October 2020, 22:27:40 UTC |
| c8a1a8d | Matthew Johnson | 16 October 2020, 22:07:02 UTC | remove extraneous test tol change | 16 October 2020, 22:07:02 UTC |
| cae5334 | Sharad Vikram | 16 October 2020, 22:04:08 UTC | fix lint error | 16 October 2020, 22:04:08 UTC |
