e075e64 | Matthew Johnson | 08 October 2020, 20:05:03 UTC | block for result in debug_nans_test | 08 October 2020, 20:05:03 UTC |
57ec9dd | jax authors | 08 October 2020, 20:01:26 UTC | Merge pull request #4483 from google:debug-nans-no-store-error PiperOrigin-RevId: 336147286 | 08 October 2020, 20:01:26 UTC |
2b76bcc | jax authors | 08 October 2020, 19:53:33 UTC | Merge pull request #4508 from zhangqiaorjc:expm_test_slow PiperOrigin-RevId: 336146064 | 08 October 2020, 19:53:33 UTC |
4dd802c | Matthew Johnson | 08 October 2020, 19:28:11 UTC | Merge branch 'master' into debug-nans-no-store-error | 08 October 2020, 19:28:11 UTC |
afef644 | Qiao Zhang | 08 October 2020, 19:09:21 UTC | Skip 2nd order grad in expm tests. | 08 October 2020, 19:09:21 UTC |
a2a6dca | jax authors | 08 October 2020, 19:06:15 UTC | Merge pull request #4507 from hawkinsp:xla PiperOrigin-RevId: 336135890 | 08 October 2020, 19:06:15 UTC |
72be97c | Peter Hawkins | 08 October 2020, 18:59:06 UTC | Update XLA. | 08 October 2020, 18:59:06 UTC |
ab48a9a | jax authors | 08 October 2020, 18:49:06 UTC | Merge pull request #4498 from jblespiau:changelist/333713171 PiperOrigin-RevId: 336132107 | 08 October 2020, 18:49:06 UTC |
145ac40 | jax authors | 08 October 2020, 17:45:01 UTC | Merge pull request #4502 from hawkinsp:linalg PiperOrigin-RevId: 336117857 | 08 October 2020, 17:45:01 UTC |
684a584 | jax authors | 08 October 2020, 17:20:34 UTC | Merge pull request #4500 from LenaMartens:changelist/336082045 PiperOrigin-RevId: 336112168 | 08 October 2020, 17:20:34 UTC |
e3d622c | Lena Martens | 08 October 2020, 14:36:05 UTC | Recast int/bool tangents to float0 in custom_jvp/vjps (also in the initial_style path). | 08 October 2020, 16:37:35 UTC |
95c0fea | jax authors | 08 October 2020, 16:06:43 UTC | Merge pull request #4501 from gnecula:tf_tests_cf1 PiperOrigin-RevId: 336096721 | 08 October 2020, 16:06:43 UTC |
2fca917 | Peter Hawkins | 08 October 2020, 14:59:52 UTC | Enable TPU linalg tests that now pass. | 08 October 2020, 14:59:52 UTC |
6fec497 | Jean-Baptiste Lespiau | 25 September 2020, 12:14:41 UTC | Add a flag to control the C++ jax.jit behavior. | 08 October 2020, 14:51:07 UTC |
dc9168b | George Necula | 08 October 2020, 09:28:42 UTC | [jax2tf] Ensure that in tests TF does not constant-fold in eager mode before compiling | 08 October 2020, 14:44:00 UTC |
8df116b | jax authors | 08 October 2020, 14:27:06 UTC | Merge pull request #4475 from bchetioui:fix_static_tests PiperOrigin-RevId: 336080718 | 08 October 2020, 14:27:06 UTC |
8299598 | Benjamin Chetioui | 08 October 2020, 13:37:20 UTC | Fix failing tests on TPU by avoiding using scalars | 08 October 2020, 13:37:20 UTC |
83060cc | jax authors | 08 October 2020, 13:06:56 UTC | Merge pull request #4491 from google:remove-test PiperOrigin-RevId: 336069850 | 08 October 2020, 13:06:56 UTC |
e65115b | jax authors | 08 October 2020, 02:55:59 UTC | Merge pull request #4489 from google:test-fixes PiperOrigin-RevId: 336003413 | 08 October 2020, 02:55:59 UTC |
30b8cc5 | Roy Frostig | 08 October 2020, 02:49:06 UTC | remove an always-skipped test, redundant by now | 08 October 2020, 02:49:06 UTC |
4b65f33 | jax authors | 08 October 2020, 01:21:27 UTC | Merge pull request #4485 from hawkinsp:nan PiperOrigin-RevId: 335993960 | 08 October 2020, 01:21:27 UTC |
1fb097b | Roy Frostig | 08 October 2020, 01:01:14 UTC | remove stale version guard in test | 08 October 2020, 01:14:58 UTC |
16cd330 | Roy Frostig | 08 October 2020, 00:59:44 UTC | skip test that is broken at unreleased (source) jaxlib | 08 October 2020, 01:14:46 UTC |
7724322 | Jean-Baptiste Lespiau | 07 October 2020, 22:35:58 UTC | Exend threading.local. The previous version could cause the following error: linear_util.py", line 244, in _most_recent_entry most_recent_entry = thread_local.most_recent_entry AttributeError: '_thread._local' object has no attribute 'most_recent_entry' PiperOrigin-RevId: 335967038 | 07 October 2020, 22:36:26 UTC |
a8b67ae | Peter Hawkins | 07 October 2020, 20:54:13 UTC | Switch implementation of jnp.isnan(x) to x != x. | 07 October 2020, 20:54:13 UTC |
9cdb139 | jax authors | 07 October 2020, 20:29:52 UTC | Merge pull request #4482 from google:tb-util-path PiperOrigin-RevId: 335940940 | 07 October 2020, 20:29:52 UTC |
4ae3352 | Jean-Baptiste Lespiau | 07 October 2020, 20:16:27 UTC | Increase the version of jaxlib, to get the tests for the C++ jit running internally. The policy is the following: - we do not create holes in the jaxlib release. - it should thus be incremented only once between releases, but may not be incremented just before the release, e.g. to trigger internally some tests (that won't run on github). PiperOrigin-RevId: 335938139 | 07 October 2020, 20:17:05 UTC |
a306c1b | jax authors | 07 October 2020, 19:51:24 UTC | Merge pull request #4479 from jblespiau:changelist/335912189 PiperOrigin-RevId: 335932313 | 07 October 2020, 19:51:24 UTC |
9b5d780 | Matthew Johnson | 07 October 2020, 19:27:11 UTC | deflake | 07 October 2020, 19:27:11 UTC |
c26c77d | Matthew Johnson | 07 October 2020, 19:17:24 UTC | fix a 'store occupied' error in jax_debug_nans This code snippet could cause a 'store occupied' error: @jit def f(x): return x + np.nan FLAGS.jax_debug_nans = True f(1) The reason is that in xla._xla_call_impl we would run a linear_util.WrappedFun twice, first via xla._xla_callable and then again directly (i.e. in op-by-op) if we got a nan on the output. Things would work fine if the second execution also raised a nan error, since then the WrappedFun wouldn't complete execution, but if the second execution does not raise an error (as in the above case, because `1 + np.nan` doesn't involve any jax primitive executions) then we'd end up with a StoreOccupied error from running the WrappedFun twice. The fix is just to intentionally allow re-running the WrappedFun, since the whole point of jax_debug_nans is to re-run functions that in normal circumstances we would only want to execute exactly once. | 07 October 2020, 19:17:24 UTC |
ab7436a | Jean-Baptiste Lespiau | 07 October 2020, 18:23:28 UTC | Ignore an incorrect mypy error. | 07 October 2020, 19:00:24 UTC |
32ac909 | Roy Frostig | 07 October 2020, 18:59:12 UTC | early return for non-existent path prefixes | 07 October 2020, 18:59:12 UTC |
c88be87 | Jean-Baptiste Lespiau | 07 October 2020, 18:14:32 UTC | Use the Python jit for the compilation in the C++ jit. Using xla_computation has been a doomed attempt, because it does not support all the features, and cannot deal with, in particular, nested tracing. Thus, we directly use the current path, and use a thread local value to access the last compiled objects from C++ (it allows to not touch the Python tracing logic). This also: - Delay the access of jax_enable_64 to after GoogleInit. PiperOrigin-RevId: 335910130 | 07 October 2020, 18:15:22 UTC |
6a4fdc9 | Benjamin Chetioui | 07 October 2020, 13:55:23 UTC | [jax2tf] Fix min_max and add_mul harnesses. Some of the array parameters were wrapped in StaticArg(), which resulted in the function not being actually converted through jax2tf. One oddity is that the test now requires "always_custom_assert" to be set to True to run properly, implying that even in compiled mode, the results produced by jax and jax2tf are inconsistent. | 07 October 2020, 13:55:23 UTC |
0290339 | jax authors | 07 October 2020, 08:31:37 UTC | Merge pull request #4468 from gnecula:tf_colab4 PiperOrigin-RevId: 335815857 | 07 October 2020, 08:31:37 UTC |
ae7e5b9 | George Necula | 07 October 2020, 08:14:35 UTC | Another attempt to save the notebook | 07 October 2020, 08:14:35 UTC |
27e5c2a | George Necula | 07 October 2020, 07:25:46 UTC | Updated the notebook with a run on a public runtime | 07 October 2020, 07:25:46 UTC |
50ce915 | George Necula | 07 October 2020, 06:46:04 UTC | Small change | 07 October 2020, 06:46:04 UTC |
ed0a098 | George Necula | 07 October 2020, 06:23:30 UTC | [jax2tf] Add the getting started notebook | 07 October 2020, 06:23:30 UTC |
2dfc5f9 | jax authors | 07 October 2020, 05:59:09 UTC | Merge pull request #4467 from google:update-pypi PiperOrigin-RevId: 335798412 | 07 October 2020, 05:59:09 UTC |
1c5a09a | Matthew Johnson | 07 October 2020, 05:23:29 UTC | update version and changelog for pypi | 07 October 2020, 05:23:29 UTC |
cf1fefc | jax authors | 06 October 2020, 20:33:20 UTC | Merge pull request #4464 from google:minor-source-sync PiperOrigin-RevId: 335709378 | 06 October 2020, 20:33:20 UTC |
5f6fab8 | Roy Frostig | 06 October 2020, 20:24:19 UTC | minor change as source sync test | 06 October 2020, 20:24:19 UTC |
6eb4f54 | jax authors | 06 October 2020, 19:45:43 UTC | Internal change PiperOrigin-RevId: 335699277 | 06 October 2020, 19:53:21 UTC |
93f4dd4 | George Necula | 06 October 2020, 11:51:51 UTC | Revert "[jax2tf] Added Getting Started notebook (#4460)" (#4463) This reverts commit 1d73952736e6142c0670879072cc99327814640e. | 06 October 2020, 11:51:51 UTC |
1d73952 | George Necula | 06 October 2020, 11:14:28 UTC | [jax2tf] Added Getting Started notebook (#4460) * [jax2tf] Added Getting Started notebook * Updated the jax2tf README * Minor cleanup | 06 October 2020, 11:14:28 UTC |
2166f0e | jax authors | 06 October 2020, 00:41:30 UTC | Merge pull request #4457 from google:tycheck-invar-eqn-context PiperOrigin-RevId: 335537630 | 06 October 2020, 00:41:30 UTC |
7efa320 | jax authors | 05 October 2020, 23:14:32 UTC | Merge pull request #4458 from hawkinsp:xla PiperOrigin-RevId: 335521726 | 05 October 2020, 23:14:32 UTC |
49d851d | Peter Hawkins | 05 October 2020, 20:14:32 UTC | Update XLA. | 05 October 2020, 20:14:32 UTC |
e797925 | Roy Frostig | 05 October 2020, 19:29:43 UTC | equation context for undefined var reads in jaxpr typechecker | 05 October 2020, 19:29:43 UTC |
91755e5 | jax authors | 05 October 2020, 19:09:35 UTC | Merge pull request #4456 from jakevdp:small-fix PiperOrigin-RevId: 335469646 | 05 October 2020, 19:09:35 UTC |
9273e8c | jax authors | 05 October 2020, 19:09:21 UTC | Merge pull request #4455 from johnpjf:patch-2 PiperOrigin-RevId: 335461715 | 05 October 2020, 19:09:21 UTC |
efeed6c | jax authors | 05 October 2020, 19:09:05 UTC | Merge pull request #4444 from jakevdp:dtype-check PiperOrigin-RevId: 335442628 | 05 October 2020, 19:09:05 UTC |
7854a23 | Jake VanderPlas | 05 October 2020, 18:50:04 UTC | remove redundant if-else | 05 October 2020, 18:50:04 UTC |
b357005 | johnpjf | 05 October 2020, 17:19:18 UTC | Update Common_Gotchas_in_JAX.ipynb Clarify that the index is clamped to the bounds of the array when accessing out of bounds. | 05 October 2020, 17:19:18 UTC |
d388b62 | jax authors | 05 October 2020, 14:17:42 UTC | Merge pull request #4442 from hawkinsp:docs PiperOrigin-RevId: 335411058 | 05 October 2020, 14:17:42 UTC |
51538bd | jax authors | 05 October 2020, 05:17:06 UTC | Merge pull request #4446 from n2cholas:master PiperOrigin-RevId: 335349341 | 05 October 2020, 05:17:06 UTC |
238b4a1 | Nicholas Vadivelu | 05 October 2020, 04:54:18 UTC | Update examples/differentially_private_sgd.py Co-authored-by: Matthew Johnson <mattjj@google.com> | 05 October 2020, 04:54:18 UTC |
0e92124 | jax authors | 03 October 2020, 10:23:55 UTC | Merge pull request #4443 from jakevdp:concrete PiperOrigin-RevId: 335180298 | 03 October 2020, 10:23:55 UTC |
2ced163 | Nicholas Vadivelu | 03 October 2020, 06:41:23 UTC | fix dp sgd example | 03 October 2020, 06:41:23 UTC |
6d5ace2 | jax authors | 03 October 2020, 01:12:39 UTC | Merge pull request #4445 from google:lax-ad-test-skip PiperOrigin-RevId: 335136639 | 03 October 2020, 01:12:39 UTC |
db4a3ee | Roy Frostig | 03 October 2020, 00:41:08 UTC | skip test that fails due to known xla:cpu bug | 03 October 2020, 00:41:08 UTC |
0328f85 | Jake VanderPlas | 02 October 2020, 22:31:02 UTC | Call check_user_dtype on all user dtypes | 03 October 2020, 00:21:44 UTC |
5cfd0e9 | jax authors | 02 October 2020, 22:56:04 UTC | Merge pull request #4441 from jakevdp:choice PiperOrigin-RevId: 335117255 | 02 October 2020, 22:56:04 UTC |
492889f | Jake VanderPlas | 02 October 2020, 22:08:21 UTC | Improve errors for failed compilations w/ core.concrete_or_error | 02 October 2020, 22:08:21 UTC |
5732428 | Peter Hawkins | 02 October 2020, 20:44:42 UTC | Fix rendering problems in lax.cond docstring. | 02 October 2020, 20:44:42 UTC |
2fccbbf | Jake VanderPlas | 02 October 2020, 20:13:21 UTC | Implement jnp.choose | 02 October 2020, 20:13:21 UTC |
95e3fcd | jax authors | 02 October 2020, 18:17:20 UTC | Merge pull request #4438 from hawkinsp:gelu PiperOrigin-RevId: 335063459 | 02 October 2020, 18:17:20 UTC |
113c806 | jax authors | 02 October 2020, 18:14:03 UTC | Merge pull request #4439 from hawkinsp:selectandscatter PiperOrigin-RevId: 335062792 | 02 October 2020, 18:14:03 UTC |
d25f404 | jax authors | 02 October 2020, 18:00:20 UTC | Merge pull request #4429 from apaszke:all_to_all_vmap_split PiperOrigin-RevId: 335059726 | 02 October 2020, 18:00:20 UTC |
2d74037 | Peter Hawkins | 02 October 2020, 14:56:08 UTC | Change select_and_scatter_add's batching rule to avoid unrolling. Add a test for select_and_scatter_add batching. | 02 October 2020, 16:02:34 UTC |
9b3bbe8 | Peter Hawkins | 02 October 2020, 13:48:07 UTC | Adds an approximate=... keyword argument to jax.nn.gelu to select between the approximate and exact formulations of gelu. Default to the approximate formulation for now. | 02 October 2020, 13:48:07 UTC |
b609040 | jax authors | 02 October 2020, 02:26:55 UTC | Merge pull request #4413 from alexminnaar:master PiperOrigin-RevId: 334946450 | 02 October 2020, 02:26:55 UTC |
9cf6159 | jax authors | 02 October 2020, 01:34:24 UTC | Merge pull request #4434 from google:numpy-dispatch-test-count PiperOrigin-RevId: 334940761 | 02 October 2020, 01:34:24 UTC |
0893b08 | Alex Minnaar | 21 September 2020, 20:59:46 UTC | histogramdd implemenation | 01 October 2020, 22:02:33 UTC |
fd39901 | Roy Frostig | 01 October 2020, 22:00:38 UTC | reduce test-case count of the numpy-dispatch CI check, to match our other x64-mode CI check | 01 October 2020, 22:00:38 UTC |
93fbae6 | jax authors | 01 October 2020, 21:54:52 UTC | Merge pull request #4431 from tomhennigan:changelist/334843453 PiperOrigin-RevId: 334902700 | 01 October 2020, 21:54:52 UTC |
f18b914 | Tom Hennigan | 01 October 2020, 17:26:07 UTC | Remove scalar float0 XLA representation. Some places in the code (e.g. lower_fun in xla.py) we need to go from an XLA shape to a ShapedArray and cannot do this if we drop shape information earlier. | 01 October 2020, 18:02:31 UTC |
9a8ee95 | jax authors | 01 October 2020, 17:39:51 UTC | Merge pull request #4419 from rsepassi:master PiperOrigin-RevId: 334846970 | 01 October 2020, 17:39:51 UTC |
a622220 | jax authors | 01 October 2020, 14:58:43 UTC | Merge pull request #4430 from LenaMartens:changelist/334791342 PiperOrigin-RevId: 334813454 | 01 October 2020, 14:58:43 UTC |
cc0114a | Lena Martens | 01 October 2020, 12:07:33 UTC | Fix dtype behavior with float0s in CustomVJP. | 01 October 2020, 14:17:51 UTC |
69fda9e | jax authors | 01 October 2020, 02:25:00 UTC | Merge pull request #4039 from LenaMartens:changelist/325216264 PiperOrigin-RevId: 334728148 | 01 October 2020, 02:25:00 UTC |
768b0c1 | jax authors | 30 September 2020, 21:44:47 UTC | Merge pull request #4415 from jakevdp:indices-concrete PiperOrigin-RevId: 334683008 | 30 September 2020, 21:44:47 UTC |
0b23a90 | Jake VanderPlas | 28 September 2020, 22:33:06 UTC | Call concrete_or_error on indices() argument | 30 September 2020, 18:44:49 UTC |
47da512 | jax authors | 30 September 2020, 17:23:34 UTC | Merge pull request #4384 from jakevdp:cov-weights PiperOrigin-RevId: 334626974 | 30 September 2020, 17:23:34 UTC |
62f1e83 | jax authors | 30 September 2020, 17:15:55 UTC | Merge pull request #4425 from akihironitta:use-raise-from PiperOrigin-RevId: 334625280 | 30 September 2020, 17:15:55 UTC |
d707ae1 | Akihiro Nitta | 30 September 2020, 15:27:03 UTC | Merge branch 'master' into use-raise-from | 30 September 2020, 15:27:03 UTC |
e61ca91 | Adam Paszke | 22 September 2020, 13:05:08 UTC | Implement split_axis for all_to_all This allows us to use `all_to_all` over a mix of vmapped and pmapped dimensions, which will be useful for `gmap`. | 30 September 2020, 08:25:36 UTC |
fa38f25 | Adam Paszke | 22 September 2020, 11:19:06 UTC | Add support for all_to_all over vmapped axes | 30 September 2020, 08:25:36 UTC |
2d94a87 | jax authors | 29 September 2020, 20:42:18 UTC | Merge pull request #4421 from jakevdp:multi-buf PiperOrigin-RevId: 334452699 | 29 September 2020, 20:42:18 UTC |
6b5a265 | jax authors | 29 September 2020, 19:07:09 UTC | Merge pull request #4414 from n2cholas:master PiperOrigin-RevId: 334432358 | 29 September 2020, 19:07:09 UTC |
40016cc | Jake VanderPlas | 29 September 2020, 18:53:17 UTC | Allow jax objects to be represented by multiple buffers | 29 September 2020, 18:53:17 UTC |
06170da | Akihiro Nitta | 29 September 2020, 16:20:00 UTC | Use `raise from` | 29 September 2020, 16:20:00 UTC |
97592c8 | Ryan Sepassi | 29 September 2020, 14:57:20 UTC | Add --exclude_test_targets to test_util | 29 September 2020, 14:57:20 UTC |
c67454d | jax authors | 29 September 2020, 13:53:17 UTC | Merge pull request #4417 from google:revert-3983-device-put-tuple PiperOrigin-RevId: 334368988 | 29 September 2020, 13:53:17 UTC |
a0e14b0 | Peter Hawkins | 29 September 2020, 13:26:11 UTC | Revert "Allow JAX objects to be represented by multiple buffers" | 29 September 2020, 13:26:11 UTC |
76e9f23 | Nicholas Vadivelu | 29 September 2020, 02:19:47 UTC | update device faq | 29 September 2020, 02:19:47 UTC |
a51a4d9 | Jake VanderPlas | 28 September 2020, 22:34:57 UTC | Add test coverage for jnp.cov aweights & fweights | 28 September 2020, 22:34:57 UTC |
ecad419 | Lena Martens | 24 September 2020, 15:29:57 UTC | Support grad with integer arguments. - Add float0 and set-up at_least_vspace to return float0 values for int/bool primals - Use Zero to wrap float0 tangents so they're correctly ignored in jvp rules - Add float0 handlers to XLA to support jit - Fix convert_element_type and tie_in jvp rules | 28 September 2020, 18:07:04 UTC |
fa11338 | jax authors | 28 September 2020, 17:00:49 UTC | Merge pull request #3983 from jakevdp:device-put-tuple PiperOrigin-RevId: 334180017 | 28 September 2020, 17:00:49 UTC |