https://github.com/google/jax

sort by:
Revision Author Date Message Commit Date
e075e64 block for result in debug_nans_test 08 October 2020, 20:05:03 UTC
57ec9dd Merge pull request #4483 from google:debug-nans-no-store-error PiperOrigin-RevId: 336147286 08 October 2020, 20:01:26 UTC
2b76bcc Merge pull request #4508 from zhangqiaorjc:expm_test_slow PiperOrigin-RevId: 336146064 08 October 2020, 19:53:33 UTC
4dd802c Merge branch 'master' into debug-nans-no-store-error 08 October 2020, 19:28:11 UTC
afef644 Skip 2nd order grad in expm tests. 08 October 2020, 19:09:21 UTC
a2a6dca Merge pull request #4507 from hawkinsp:xla PiperOrigin-RevId: 336135890 08 October 2020, 19:06:15 UTC
72be97c Update XLA. 08 October 2020, 18:59:06 UTC
ab48a9a Merge pull request #4498 from jblespiau:changelist/333713171 PiperOrigin-RevId: 336132107 08 October 2020, 18:49:06 UTC
145ac40 Merge pull request #4502 from hawkinsp:linalg PiperOrigin-RevId: 336117857 08 October 2020, 17:45:01 UTC
684a584 Merge pull request #4500 from LenaMartens:changelist/336082045 PiperOrigin-RevId: 336112168 08 October 2020, 17:20:34 UTC
e3d622c Recast int/bool tangents to float0 in custom_jvp/vjps (also in the initial_style path). 08 October 2020, 16:37:35 UTC
95c0fea Merge pull request #4501 from gnecula:tf_tests_cf1 PiperOrigin-RevId: 336096721 08 October 2020, 16:06:43 UTC
2fca917 Enable TPU linalg tests that now pass. 08 October 2020, 14:59:52 UTC
6fec497 Add a flag to control the C++ jax.jit behavior. 08 October 2020, 14:51:07 UTC
dc9168b [jax2tf] Ensure that in tests TF does not constant-fold in eager mode before compiling 08 October 2020, 14:44:00 UTC
8df116b Merge pull request #4475 from bchetioui:fix_static_tests PiperOrigin-RevId: 336080718 08 October 2020, 14:27:06 UTC
8299598 Fix failing tests on TPU by avoiding using scalars 08 October 2020, 13:37:20 UTC
83060cc Merge pull request #4491 from google:remove-test PiperOrigin-RevId: 336069850 08 October 2020, 13:06:56 UTC
e65115b Merge pull request #4489 from google:test-fixes PiperOrigin-RevId: 336003413 08 October 2020, 02:55:59 UTC
30b8cc5 remove an always-skipped test, redundant by now 08 October 2020, 02:49:06 UTC
4b65f33 Merge pull request #4485 from hawkinsp:nan PiperOrigin-RevId: 335993960 08 October 2020, 01:21:27 UTC
1fb097b remove stale version guard in test 08 October 2020, 01:14:58 UTC
16cd330 skip test that is broken at unreleased (source) jaxlib 08 October 2020, 01:14:46 UTC
7724322 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 Switch implementation of jnp.isnan(x) to x != x. 07 October 2020, 20:54:13 UTC
9cdb139 Merge pull request #4482 from google:tb-util-path PiperOrigin-RevId: 335940940 07 October 2020, 20:29:52 UTC
4ae3352 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 Merge pull request #4479 from jblespiau:changelist/335912189 PiperOrigin-RevId: 335932313 07 October 2020, 19:51:24 UTC
9b5d780 deflake 07 October 2020, 19:27:11 UTC
c26c77d 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 Ignore an incorrect mypy error. 07 October 2020, 19:00:24 UTC
32ac909 early return for non-existent path prefixes 07 October 2020, 18:59:12 UTC
c88be87 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 [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 Merge pull request #4468 from gnecula:tf_colab4 PiperOrigin-RevId: 335815857 07 October 2020, 08:31:37 UTC
ae7e5b9 Another attempt to save the notebook 07 October 2020, 08:14:35 UTC
27e5c2a Updated the notebook with a run on a public runtime 07 October 2020, 07:25:46 UTC
50ce915 Small change 07 October 2020, 06:46:04 UTC
ed0a098 [jax2tf] Add the getting started notebook 07 October 2020, 06:23:30 UTC
2dfc5f9 Merge pull request #4467 from google:update-pypi PiperOrigin-RevId: 335798412 07 October 2020, 05:59:09 UTC
1c5a09a update version and changelog for pypi 07 October 2020, 05:23:29 UTC
cf1fefc Merge pull request #4464 from google:minor-source-sync PiperOrigin-RevId: 335709378 06 October 2020, 20:33:20 UTC
5f6fab8 minor change as source sync test 06 October 2020, 20:24:19 UTC
6eb4f54 Internal change PiperOrigin-RevId: 335699277 06 October 2020, 19:53:21 UTC
93f4dd4 Revert "[jax2tf] Added Getting Started notebook (#4460)" (#4463) This reverts commit 1d73952736e6142c0670879072cc99327814640e. 06 October 2020, 11:51:51 UTC
1d73952 [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 Merge pull request #4457 from google:tycheck-invar-eqn-context PiperOrigin-RevId: 335537630 06 October 2020, 00:41:30 UTC
7efa320 Merge pull request #4458 from hawkinsp:xla PiperOrigin-RevId: 335521726 05 October 2020, 23:14:32 UTC
49d851d Update XLA. 05 October 2020, 20:14:32 UTC
e797925 equation context for undefined var reads in jaxpr typechecker 05 October 2020, 19:29:43 UTC
91755e5 Merge pull request #4456 from jakevdp:small-fix PiperOrigin-RevId: 335469646 05 October 2020, 19:09:35 UTC
9273e8c Merge pull request #4455 from johnpjf:patch-2 PiperOrigin-RevId: 335461715 05 October 2020, 19:09:21 UTC
efeed6c Merge pull request #4444 from jakevdp:dtype-check PiperOrigin-RevId: 335442628 05 October 2020, 19:09:05 UTC
7854a23 remove redundant if-else 05 October 2020, 18:50:04 UTC
b357005 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 Merge pull request #4442 from hawkinsp:docs PiperOrigin-RevId: 335411058 05 October 2020, 14:17:42 UTC
51538bd Merge pull request #4446 from n2cholas:master PiperOrigin-RevId: 335349341 05 October 2020, 05:17:06 UTC
238b4a1 Update examples/differentially_private_sgd.py Co-authored-by: Matthew Johnson <mattjj@google.com> 05 October 2020, 04:54:18 UTC
0e92124 Merge pull request #4443 from jakevdp:concrete PiperOrigin-RevId: 335180298 03 October 2020, 10:23:55 UTC
2ced163 fix dp sgd example 03 October 2020, 06:41:23 UTC
6d5ace2 Merge pull request #4445 from google:lax-ad-test-skip PiperOrigin-RevId: 335136639 03 October 2020, 01:12:39 UTC
db4a3ee skip test that fails due to known xla:cpu bug 03 October 2020, 00:41:08 UTC
0328f85 Call check_user_dtype on all user dtypes 03 October 2020, 00:21:44 UTC
5cfd0e9 Merge pull request #4441 from jakevdp:choice PiperOrigin-RevId: 335117255 02 October 2020, 22:56:04 UTC
492889f Improve errors for failed compilations w/ core.concrete_or_error 02 October 2020, 22:08:21 UTC
5732428 Fix rendering problems in lax.cond docstring. 02 October 2020, 20:44:42 UTC
2fccbbf Implement jnp.choose 02 October 2020, 20:13:21 UTC
95e3fcd Merge pull request #4438 from hawkinsp:gelu PiperOrigin-RevId: 335063459 02 October 2020, 18:17:20 UTC
113c806 Merge pull request #4439 from hawkinsp:selectandscatter PiperOrigin-RevId: 335062792 02 October 2020, 18:14:03 UTC
d25f404 Merge pull request #4429 from apaszke:all_to_all_vmap_split PiperOrigin-RevId: 335059726 02 October 2020, 18:00:20 UTC
2d74037 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 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 Merge pull request #4413 from alexminnaar:master PiperOrigin-RevId: 334946450 02 October 2020, 02:26:55 UTC
9cf6159 Merge pull request #4434 from google:numpy-dispatch-test-count PiperOrigin-RevId: 334940761 02 October 2020, 01:34:24 UTC
0893b08 histogramdd implemenation 01 October 2020, 22:02:33 UTC
fd39901 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 Merge pull request #4431 from tomhennigan:changelist/334843453 PiperOrigin-RevId: 334902700 01 October 2020, 21:54:52 UTC
f18b914 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 Merge pull request #4419 from rsepassi:master PiperOrigin-RevId: 334846970 01 October 2020, 17:39:51 UTC
a622220 Merge pull request #4430 from LenaMartens:changelist/334791342 PiperOrigin-RevId: 334813454 01 October 2020, 14:58:43 UTC
cc0114a Fix dtype behavior with float0s in CustomVJP. 01 October 2020, 14:17:51 UTC
69fda9e Merge pull request #4039 from LenaMartens:changelist/325216264 PiperOrigin-RevId: 334728148 01 October 2020, 02:25:00 UTC
768b0c1 Merge pull request #4415 from jakevdp:indices-concrete PiperOrigin-RevId: 334683008 30 September 2020, 21:44:47 UTC
0b23a90 Call concrete_or_error on indices() argument 30 September 2020, 18:44:49 UTC
47da512 Merge pull request #4384 from jakevdp:cov-weights PiperOrigin-RevId: 334626974 30 September 2020, 17:23:34 UTC
62f1e83 Merge pull request #4425 from akihironitta:use-raise-from PiperOrigin-RevId: 334625280 30 September 2020, 17:15:55 UTC
d707ae1 Merge branch 'master' into use-raise-from 30 September 2020, 15:27:03 UTC
e61ca91 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 Add support for all_to_all over vmapped axes 30 September 2020, 08:25:36 UTC
2d94a87 Merge pull request #4421 from jakevdp:multi-buf PiperOrigin-RevId: 334452699 29 September 2020, 20:42:18 UTC
6b5a265 Merge pull request #4414 from n2cholas:master PiperOrigin-RevId: 334432358 29 September 2020, 19:07:09 UTC
40016cc Allow jax objects to be represented by multiple buffers 29 September 2020, 18:53:17 UTC
06170da Use `raise from` 29 September 2020, 16:20:00 UTC
97592c8 Add --exclude_test_targets to test_util 29 September 2020, 14:57:20 UTC
c67454d Merge pull request #4417 from google:revert-3983-device-put-tuple PiperOrigin-RevId: 334368988 29 September 2020, 13:53:17 UTC
a0e14b0 Revert "Allow JAX objects to be represented by multiple buffers" 29 September 2020, 13:26:11 UTC
76e9f23 update device faq 29 September 2020, 02:19:47 UTC
a51a4d9 Add test coverage for jnp.cov aweights & fweights 28 September 2020, 22:34:57 UTC
ecad419 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 Merge pull request #3983 from jakevdp:device-put-tuple PiperOrigin-RevId: 334180017 28 September 2020, 17:00:49 UTC
back to top