| bd1708c | Peter Hawkins | 31 March 2020, 14:02:38 UTC | Update changelog and README for jaxlib 0.1.43. (#2556) | 31 March 2020, 14:02:38 UTC |
| 59ed4ae | George Necula | 31 March 2020, 09:11:47 UTC | Disable test_while_loop_of_pmap on all platforms Issue: #2554 Disable the test so that we can continue the google3 tests | 31 March 2020, 09:12:35 UTC |
| fd52fbf | George Necula | 31 March 2020, 08:36:47 UTC | Fix import in benchmarks This works on my machine as 'python benchmarks/pmap_benchmark.py'. It also follows the code in examples. This will need a copybara rule to change the import to 'from jax.benchmarks import benchmark' | 31 March 2020, 08:48:08 UTC |
| a4ceae1 | Matthew Johnson | 31 March 2020, 05:12:38 UTC | fix link in custom derivatives tutorial notebook | 31 March 2020, 05:12:38 UTC |
| 27604c3 | Matthew Johnson | 31 March 2020, 05:11:35 UTC | fix typo in notebook | 31 March 2020, 05:11:35 UTC |
| e017a92 | Matthew Johnson | 31 March 2020, 05:06:00 UTC | fix typo | 31 March 2020, 05:06:00 UTC |
| 7a9c550 | Matthew Johnson | 31 March 2020, 04:30:47 UTC | add sphinx-autodoc-typehints to travis install | 31 March 2020, 04:30:47 UTC |
| e7be43d | Matthew Johnson | 31 March 2020, 04:09:12 UTC | update api.py docstrings for sphinx highlighting | 31 March 2020, 04:09:12 UTC |
| cda505a | Matthew Johnson | 31 March 2020, 03:46:09 UTC | Merge pull request #2551 from google/enable-mpc-tests try re-enabling control tests that trigger #2507 | 31 March 2020, 03:46:09 UTC |
| 909fee6 | Matthew Johnson | 31 March 2020, 03:22:04 UTC | try adding sphinx-autodoc-typehints | 31 March 2020, 03:22:04 UTC |
| b015e57 | Matthew Johnson | 31 March 2020, 03:12:33 UTC | try re-enabling control tests that trigger #2507 | 31 March 2020, 03:12:33 UTC |
| 1f03d48 | Matthew Johnson | 31 March 2020, 03:10:39 UTC | try resetting global tracer state in loops_test.py attempting to address #2507 | 31 March 2020, 03:10:39 UTC |
| de37eae | Matthew Johnson | 31 March 2020, 03:10:16 UTC | Merge pull request #2550 from google/update-custom-derivatives-notebook update custom derivatives tutorial notebook | 31 March 2020, 03:10:16 UTC |
| 15009c9 | Matthew Johnson | 31 March 2020, 02:45:45 UTC | add docstring for defjvps, fix sphinx docs | 31 March 2020, 02:45:45 UTC |
| bd726fc | Matthew Johnson | 31 March 2020, 02:37:11 UTC | update custom derivatives tutorial notebook * add clip_gradient example * add defjvps convenience wrapper | 31 March 2020, 02:37:11 UTC |
| b7d2cbd | Matthew Johnson | 31 March 2020, 01:34:50 UTC | Merge pull request #2548 from google/while-loop-of-pmap-bug fix a while-loop-of-pmap bug (thanks @jaspersnoek) | 31 March 2020, 01:34:50 UTC |
| 375575a | Matthew Johnson | 31 March 2020, 00:53:47 UTC | skip new test on cpu unless num_devices > 1 | 31 March 2020, 00:53:47 UTC |
| df3b5fe | Matthew Johnson | 31 March 2020, 00:48:07 UTC | fix a while-loop-of-pmap bug (thanks @jaspersnoek) | 31 March 2020, 00:48:07 UTC |
| acb2414 | Peter Hawkins | 30 March 2020, 23:43:23 UTC | Update XLA to fix CUDA 9.2 build problem. (#2547) | 30 March 2020, 23:43:23 UTC |
| 0c7e134 | Peter Hawkins | 30 March 2020, 20:56:54 UTC | Update XLA and increment jaxlib version to 0.1.43. (#2546) | 30 March 2020, 20:56:54 UTC |
| 1275bc5 | Matthew Johnson | 30 March 2020, 20:50:32 UTC | don't pass `-n 1` to pytest in travis | 30 March 2020, 20:50:32 UTC |
| b43051e | Matthew Johnson | 30 March 2020, 20:49:56 UTC | minor fix to custom_transforms | 30 March 2020, 20:49:56 UTC |
| 3d0c187 | Matthew Johnson | 30 March 2020, 20:38:05 UTC | Merge pull request #2543 from google/process-custom-jvp-default-implementation comments/defaults for process_custom_{jv,vj}p_call | 30 March 2020, 20:38:05 UTC |
| ea9401d | Matthew Johnson | 30 March 2020, 19:47:13 UTC | Merge pull request #2541 from google/xla-computation-duck-typing allow duck-typing in xla_computation arguments | 30 March 2020, 19:47:13 UTC |
| 70a3f47 | Matthew Johnson | 30 March 2020, 18:57:03 UTC | comments/defaults for process_custom_{jv,vj}p_call | 30 March 2020, 19:02:25 UTC |
| f766c5e | Matthew Johnson | 30 March 2020, 18:31:29 UTC | allow duck-typing in xla_computation arguments | 30 March 2020, 18:31:29 UTC |
| 9d8823c | Matthew Johnson | 30 March 2020, 07:41:04 UTC | add initial_style_staging to custom_transforms | 30 March 2020, 07:41:04 UTC |
| a6a837a | Matthew Johnson | 30 March 2020, 07:35:45 UTC | add some stage_out=True indicators | 30 March 2020, 07:35:45 UTC |
| fab6600 | Matthew Johnson | 30 March 2020, 06:31:07 UTC | Merge pull request #2539 from google/remat-fix remove remat context check, add initial staging | 30 March 2020, 06:31:07 UTC |
| bdc0c3b | Matthew Johnson | 30 March 2020, 06:29:55 UTC | remove remat context check, add initial staging | 30 March 2020, 06:29:55 UTC |
| e7f8503 | Matthew Johnson | 30 March 2020, 06:04:00 UTC | Merge pull request #2500 from google/custom-jvp-fix revise custom_jvp / custom_vjp rule jaxpr staging | 30 March 2020, 06:04:00 UTC |
| 74d358d | Matthew Johnson | 30 March 2020, 06:00:40 UTC | skip ode test on import error (internal) | 30 March 2020, 06:00:40 UTC |
| b2b68e5 | Matthew Johnson | 30 March 2020, 03:51:51 UTC | fix bugs, add tests | 30 March 2020, 05:53:37 UTC |
| 7a4c4d5 | Matthew Johnson | 30 March 2020, 03:48:08 UTC | use custom_jvp for internal functions | 30 March 2020, 03:48:08 UTC |
| 6193e5e | Matthew Johnson | 28 March 2020, 21:15:46 UTC | revamp custom_jvp/vjp implementation to fix bugs Co-authored-by: Dougal Maclaurin <dougalm@google.com> | 30 March 2020, 02:35:01 UTC |
| 954334c | Matthew Johnson | 30 March 2020, 01:05:48 UTC | Merge pull request #2537 from duvenaud/morejets Adds lots of trivial jet rules, partly addresses #2431 | 30 March 2020, 01:05:48 UTC |
| 0f11f9c | Matthew Johnson | 30 March 2020, 00:02:19 UTC | Merge pull request #2538 from google/pfau workaround for pmap output PRED arrays on cpu/gpu | 30 March 2020, 00:02:19 UTC |
| 1762a86 | Matthew Johnson | 29 March 2020, 21:45:17 UTC | workaround for pmap output PRED arrays on cpu/gpu | 29 March 2020, 21:45:17 UTC |
| 305dd8c | Matthew Johnson | 29 March 2020, 21:43:42 UTC | Merge pull request #2536 from google/issue2534 add docstring / reference doc link for axis_index | 29 March 2020, 21:43:42 UTC |
| fcc1e76 | Matthew Johnson | 29 March 2020, 20:56:26 UTC | add docstring / reference doc link for axis_index fixes #2534 | 29 March 2020, 20:56:26 UTC |
| 614d39d | Matthew Johnson | 29 March 2020, 20:37:49 UTC | Merge pull request #2527 from lucasb-eyer/patch-1 Make it more explicit that default JVP assumes |R | 29 March 2020, 20:37:49 UTC |
| ead8011 | David Duvenaud | 29 March 2020, 20:28:17 UTC | Added lots of trivial jet rules. Co-Authored-By: jessebett <jessebett@gmail.com> Co-Authored-By: Jacob Kelly <jacob.kelly@mail.utoronto.ca> | 29 March 2020, 20:28:17 UTC |
| 67283a0 | Matthew Johnson | 28 March 2020, 20:52:40 UTC | add new custom_jvp tests from #2500 Co-authored-by: Dougal Maclaurin <dougalm@google.com> | 28 March 2020, 22:33:20 UTC |
| bcc5191 | Matthew Johnson | 28 March 2020, 22:15:47 UTC | Merge pull request #2533 from google/core-trace-type-annotations add type annotations to core.py tracing machinery | 28 March 2020, 22:15:47 UTC |
| f99720b | Matthew Johnson | 28 March 2020, 21:55:58 UTC | add type annotations to core.py tracing machinery also add .copy() method to core.trace_state global trace state | 28 March 2020, 21:58:35 UTC |
| cdbcbb0 | Matthew Johnson | 28 March 2020, 19:31:40 UTC | Merge pull request #2532 from google/sharded-device-array-handlers add ShardedDeviceArray to ad vspace op handlers | 28 March 2020, 19:31:40 UTC |
| 1b59789 | Matthew Johnson | 28 March 2020, 18:56:12 UTC | add ShardedDeviceArray to ad vspace op handlers fixes #2529 (thanks, @dpfau !) | 28 March 2020, 18:56:12 UTC |
| 415cde5 | Lucas Beyer | 28 March 2020, 11:32:44 UTC | Make it more explicit that default JVP assumes |R It's just an attempt to make this implicit assumption, as it only became clear to me after our discussion in chat, not after reading this. | 28 March 2020, 11:32:44 UTC |
| f371bfc | Peter Hawkins | 28 March 2020, 01:24:26 UTC | Improve speed of LU decomposition on TPU. (#2526) Increase the block size, which helps with compilation time. Merge the two row permutations in the outer loop, which means we do row-at-a-time gathers. | 28 March 2020, 01:24:26 UTC |
| 6d0810a | Matthew Johnson | 27 March 2020, 19:46:01 UTC | Merge pull request #2501 from botev/isclose Making isclose handle correctly infinite and NaN values. | 27 March 2020, 19:46:01 UTC |
| 24bbd2b | Skye Wanderman-Milne | 27 March 2020, 17:50:57 UTC | Fix pmap_benchmark.py import (#2524) | 27 March 2020, 17:50:57 UTC |
| db0ac98 | Matthew Johnson | 27 March 2020, 16:35:07 UTC | Merge pull request #2523 from sharadmv/tangent-fix Instantiate zeros in _custom_vjp_call_jaxpr_jvp | 27 March 2020, 16:35:07 UTC |
| 0499b84 | Sharad Vikram | 27 March 2020, 02:07:48 UTC | Instantiate zeros in _custom_vjp_call_jaxpr_jvp | 27 March 2020, 02:07:48 UTC |
| 42dbfd4 | Matthew Johnson | 26 March 2020, 23:52:29 UTC | attempt to fix link formatting with nbsphinx | 26 March 2020, 23:52:29 UTC |
| a4215f2 | Matthew Johnson | 26 March 2020, 04:01:00 UTC | Merge pull request #2513 from google/temp-revert-2026 temporarily revert parts of #2026 pending bug fix | 26 March 2020, 04:01:00 UTC |
| 3700fab | Matthew Johnson | 26 March 2020, 03:24:11 UTC | remove deprecation warnings | 26 March 2020, 03:24:11 UTC |
| c3e3d48 | Matthew Johnson | 26 March 2020, 03:19:49 UTC | temporarily revert parts of #2026 pending bug fix | 26 March 2020, 03:19:49 UTC |
| 3274747 | Matthew Johnson | 26 March 2020, 01:17:55 UTC | fix derivatives reference (wrong Rudin!) | 26 March 2020, 01:17:55 UTC |
| fc0f875 | Matthew Johnson | 26 March 2020, 00:05:57 UTC | improve ref to Tao's 3rd edition of Analysis I | 26 March 2020, 00:05:57 UTC |
| da9b523 | Matthew Johnson | 25 March 2020, 21:53:23 UTC | remove incorrect sentence in notebook | 25 March 2020, 21:53:23 UTC |
| 9ec0ac6 | botev | 25 March 2020, 09:59:43 UTC | Added test for isclose. | 25 March 2020, 09:59:43 UTC |
| f88d49b | George Necula | 24 March 2020, 09:22:49 UTC | Added FAQ entry about creating JAX arrays | 25 March 2020, 09:55:24 UTC |
| 86e3046 | George Necula | 22 March 2020, 05:47:14 UTC | Added a FAQ for impure functions | 25 March 2020, 09:55:24 UTC |
| 6f2f779 | George Necula | 19 March 2020, 13:55:16 UTC | Started a FAQ for JAX | 25 March 2020, 09:55:24 UTC |
| d1e8f43 | George Necula | 25 March 2020, 09:07:50 UTC | Better error message for indexing with floats (#2496) Now the error message is 'Indexer must have integer or boolean type'. Before it was 'len() of unsized object' | 25 March 2020, 09:07:50 UTC |
| a4922ce | Matthew Johnson | 25 March 2020, 05:30:43 UTC | Merge pull request #2503 from google/issue2502 improve custom_jvp error messages | 25 March 2020, 05:30:43 UTC |
| 74c2050 | Matthew Johnson | 25 March 2020, 03:43:33 UTC | improve custom_jvp error messages, fixes #2502 | 25 March 2020, 04:45:50 UTC |
| b05ac57 | botev | 25 March 2020, 02:10:06 UTC | Making isclose handle correctly infinite and NaN values. | 25 March 2020, 02:13:57 UTC |
| 0cf84f9 | Matthew Johnson | 24 March 2020, 01:43:02 UTC | fix custom_transforms bug | 24 March 2020, 01:43:02 UTC |
| 8c4d5ce | Matthew Johnson | 24 March 2020, 00:42:40 UTC | Merge pull request #2493 from google/custom-transforms-cleanup separate out deprecated custom_transforms stuff | 24 March 2020, 00:42:40 UTC |
| 83bf048 | Matthew Johnson | 23 March 2020, 21:29:22 UTC | separate out deprecated custom_transforms stuff | 23 March 2020, 23:45:28 UTC |
| c76f32b | Matthew Johnson | 23 March 2020, 19:18:59 UTC | remove jarrett and _make_graphviz, bitrot might want to revive jarrett later! | 23 March 2020, 19:18:59 UTC |
| 61d885c | George Necula | 23 March 2020, 06:08:16 UTC | Revert previous commit (remove old custom_transforms API) (#2489) * Revert "Add back support for custom_transforms" This reverts commit 3abe61f1aed3810f7dc922c87e5cfec1c879fd27. | 23 March 2020, 06:08:16 UTC |
| f658eb5 | George Necula | 22 March 2020, 18:50:06 UTC | Add back support for custom_transforms (#2484) * add also the tests * mark the old APIs as deprecated | 22 March 2020, 18:50:06 UTC |
| 069cb3e | Matthew Johnson | 22 March 2020, 05:57:51 UTC | Merge pull request #2026 from google/custom-transforms3 add custom_jvp / vjp, delete custom_transforms | 22 March 2020, 05:57:51 UTC |
| 7e480fa | Matthew Johnson | 15 January 2020, 23:00:38 UTC | add custom_jvp / vjp, delete custom_transforms | 22 March 2020, 05:08:03 UTC |
| 6876271 | Matthew Johnson | 22 March 2020, 04:20:12 UTC | bump tolerance for mvn logpdf x64 test | 22 March 2020, 04:20:12 UTC |
| 15c8d4c | Matthew Johnson | 22 March 2020, 03:57:25 UTC | update version and changelog for pypi | 22 March 2020, 03:57:25 UTC |
| a669eda | Matthew Johnson | 22 March 2020, 02:56:03 UTC | Merge pull request #2470 from GeorgOstrovski/changelist/301855313 Make jax.numpy.squeeze as strict as numpy.squeeze about axis shape. | 22 March 2020, 02:56:03 UTC |
| 7b0ee9a | Matthew Johnson | 21 March 2020, 22:42:59 UTC | improve implementation of MVN logpdf (#2481) fixes #2314 I also added a bit more test coverage, but not a ton: scipy has different batch shape semantics and default arguments than I might expect, so I didn't bother to implement those (and left some test cases commented out). I ran into this surprising scipy bug: ```python In [1]: from scipy.stats import multivariate_normal In [2]: import numpy as np In [3]: args = [np.array(1., np.float32), np.array(2., np.float64), np.array(3., np.float64)] In [4]: print([x.shape for x in args]) [(), (), ()] In [5]: multivariate_normal.logpdf(*args) Out[5]: -1.6349113442053944 In [6]: print([x.shape for x in args]) [(), (1,), (1, 1)] ``` Mutated arguments! But it depends on dtype promotion: ```python In [7]: args = [np.array(1., np.float32), np.array(2., np.float32), np.array(3., np.float32)] In [8]: print([x.shape for x in args]) [(), (), ()] In [9]: multivariate_normal.logpdf(*args) Out[9]: -1.6349113442053944 In [10]: print([x.shape for x in args]) [(), (), ()] ``` | 21 March 2020, 22:42:59 UTC |
| e456edf | Matthew Johnson | 21 March 2020, 19:22:44 UTC | Merge pull request #2480 from google/solve-triangular-vectors make lax_linalg.solve_triangular allow vector rhs | 21 March 2020, 19:22:44 UTC |
| 93d3e34 | Matthew Johnson | 21 March 2020, 17:46:07 UTC | make lax_linalg.solve_triangular allow vector rhs also add tests for jax.scipy.linalg.cho_solve | 21 March 2020, 17:46:07 UTC |
| fcdbe63 | George Necula | 21 March 2020, 16:38:46 UTC | Trigger a Travis build (#2477) * Remove more unused imports * Fix warnings in travis.yml | 21 March 2020, 16:38:46 UTC |
| e41a5ea | George Necula | 21 March 2020, 14:45:59 UTC | Pytype2 (#2476) * Try to install pytype from pip | 21 March 2020, 14:45:59 UTC |
| 428377a | George Necula | 21 March 2020, 12:54:30 UTC | Added type annotations and removed unused imports (#2472) * Added type annotations and removed unused imports * Adjusted type hints for pytype | 21 March 2020, 12:54:30 UTC |
| 9331fc5 | George Necula | 21 March 2020, 12:53:35 UTC | Added pytype checking to Travis (#2475) | 21 March 2020, 12:53:35 UTC |
| f7e6b26 | Georg Ostrovski | 19 March 2020, 18:17:02 UTC | Make jax.numpy.squeeze as strict as numpy.squeeze about axis shape. Raise error if an axis explicitly selected to be squeezed has shape != 1. | 20 March 2020, 13:46:06 UTC |
| 578e5cf | Peter Hawkins | 19 March 2020, 18:54:04 UTC | Fix return type for vjp. (#2462) Fix vjp doc string. | 19 March 2020, 18:54:04 UTC |
| 592ba02 | Peter Hawkins | 19 March 2020, 18:53:24 UTC | Fix nondeterminism in test ordering in Python 3.6. (#2460) Set orders aren't deterministic, and this made pytest-xdist complain. | 19 March 2020, 18:53:24 UTC |
| 7f8ce8f | Matthew Johnson | 19 March 2020, 18:28:35 UTC | fix test errors from previous commit | 19 March 2020, 18:33:00 UTC |
| 1d0b7e2 | Matthew Johnson | 19 March 2020, 18:26:29 UTC | make jaxpr pretty-print show multiple outputs | 19 March 2020, 18:26:29 UTC |
| c3f8909 | Peter Hawkins | 19 March 2020, 16:41:19 UTC | Update jaxlib version in README.md. (#2461) | 19 March 2020, 16:41:19 UTC |
| d11a9ab | Trevor Cai | 19 March 2020, 15:35:00 UTC | Expose jax.lax.all_gather (#2449) * Expose jax.lax.all_gather * add all_gather to RTD | 19 March 2020, 15:35:00 UTC |
| c7f211d | Srinivas Vasudevan | 19 March 2020, 14:29:37 UTC | Update JAX to use XLA hyperbolic functions. (#2415) | 19 March 2020, 14:29:37 UTC |
| afdd1a7 | Peter Hawkins | 19 March 2020, 14:28:29 UTC | Add more return types to api.py. (#2452) | 19 March 2020, 14:28:29 UTC |
| cecfb37 | Peter Hawkins | 19 March 2020, 13:57:11 UTC | Increment jaxlib version to 0.1.42. (#2457) Update XLA. | 19 March 2020, 13:57:11 UTC |
| 78c1f6b | George Necula | 19 March 2020, 07:54:37 UTC | Increased tolerance for testScipySpecialFun (#2454) Prevent failures on TPU | 19 March 2020, 07:54:37 UTC |
| cd7ab0a | George Necula | 19 March 2020, 05:56:59 UTC | Changed to pmap_benchmark to make it runnable in Google (#2448) | 19 March 2020, 05:56:59 UTC |
| 2998a21 | George Necula | 19 March 2020, 05:55:43 UTC | Updated Common Gotchas (#2435) * Minor update to docs; trigger readthedocs * Updated Common Gotchas notebook Handle errors explicitly, otherwise it is too hard to test the notebook by 'Run all' * Added a section about pure functions to Common Gotchas | 19 March 2020, 05:55:43 UTC |
| 68b32bf | Peter Hawkins | 18 March 2020, 21:06:05 UTC | Add mypy type checking (#2430) * Add type annotations to make mypy pass. * Add mypy to .travis.yml. | 18 March 2020, 21:06:05 UTC |
