https://github.com/google/jax

sort by:
Revision Author Date Message Commit Date
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
c4a08af document setting default device and array creation 27 September 2020, 00:36:40 UTC
5b3cbc5 Merge pull request #4342 from google:improve-tracer-error PiperOrigin-RevId: 333841912 26 September 2020, 00:44:21 UTC
23a25da fix test 25 September 2020, 23:02:32 UTC
aaa5724 Merge pull request #4412 from jakevdp:axis-concrete PiperOrigin-RevId: 333822998 25 September 2020, 22:42:58 UTC
61fcea9 tweak error to mention *flattened* positions 25 September 2020, 22:35:44 UTC
f66c2ee Merge branch 'master' into improve-tracer-error 25 September 2020, 22:13:14 UTC
185590f Use core.concrete_or_error() to improve errors in reductions 25 September 2020, 21:18:46 UTC
4d55683 Merge pull request #4411 from google:omnistaging-doc-fix-links PiperOrigin-RevId: 333787009 25 September 2020, 19:35:24 UTC
be6e85a fix links in omnistaging doc (hopefully...) 25 September 2020, 19:10:03 UTC
5e19481 Merge pull request #4410 from google:add-omnistaging-doc PiperOrigin-RevId: 333779252 25 September 2020, 18:59:22 UTC
90b272a add author line, fix typo 25 September 2020, 18:50:32 UTC
0fb12a2 add omnistaging doc 25 September 2020, 18:48:49 UTC
d1f8022 Allow jax objects to be represented by multiple buffers 25 September 2020, 18:09:08 UTC
8d24d0a Merge pull request #4368 from jakevdp:unop-check-array PiperOrigin-RevId: 333765605 25 September 2020, 17:56:00 UTC
95389e8 Merge pull request #4406 from gnecula:print_debug PiperOrigin-RevId: 333727118 25 September 2020, 14:12:58 UTC
f46a422 [host_callback] improved logging, and add instructions for debugging 25 September 2020, 12:28:23 UTC
ca32966 Merge pull request #4403 from SIben:eigh-edge-case PiperOrigin-RevId: 333679325 25 September 2020, 07:11:36 UTC
a1da363 Merge pull request #4402 from gnecula:omni_print PiperOrigin-RevId: 333676897 25 September 2020, 06:47:00 UTC
fd5294b [jax2tf] Fix casting in translation of eigh for empty arrays. In the case when the operand is complex, we still need to cast the resulting empty eigenvalues to the corresponding floating point type. 24 September 2020, 14:49:06 UTC
1bf8036 Merge pull request #4170 from apaszke:all_to_all PiperOrigin-RevId: 333502317 24 September 2020, 13:31:46 UTC
acd4cc5 Allow vmapping all_to_all and implement a (slow) CPU and GPU translation This allows pmapping vmapped computations that use `all_to_all` or `pswapaxes` inside. It also includes a very slow CPU and GPU translation rule that might be useful for debugging programs locally, since XLA only implements the `AllToAll` collective on TPUs. Fixes #4141. 24 September 2020, 12:18:59 UTC
03970e9 Tag tests that require omnistaging 24 September 2020, 12:08:07 UTC
2b61513 [host_callback] Fixes and documentation changes for omnistaging As a nice benefit from omnistaging, the id_print are executed in program order even if the result is not ued. We keep the functions returning values for backwards compatibility, but also so the can be used inside expressions and also so that they can tap values from the backward computation. Omnistaging introduces jaxpr with no invars, which broke host_callback. 24 September 2020, 11:24:02 UTC
f7a09c6 Merge pull request #4388 from apaszke:all_to_all_translation PiperOrigin-RevId: 333480692 24 September 2020, 10:28:32 UTC
b8396c1 Merge pull request #4364 from botev:api PiperOrigin-RevId: 333478420 24 September 2020, 10:09:12 UTC
aa85166 Merge pull request #4387 from gnecula:print_doc PiperOrigin-RevId: 333451138 24 September 2020, 06:11:48 UTC
20c8c1a Merge pull request #4400 from google:issue4389 PiperOrigin-RevId: 333448765 24 September 2020, 05:50:49 UTC
bdcefbd Merge pull request #4401 from google:update-pypi PiperOrigin-RevId: 333440474 24 September 2020, 04:49:39 UTC
951bc75 update version and changelog for pypi 24 September 2020, 04:26:10 UTC
ebf7c1b add jax logo file 24 September 2020, 04:09:44 UTC
c7e0ef4 Merge pull request #4398 from google:lift-randomness-limit PiperOrigin-RevId: 333433816 24 September 2020, 03:51:27 UTC
d607164 make_jaxpr return_shape use ShapeDtypeStruct, test 24 September 2020, 03:41:57 UTC
71f5f99 skip checks in big randomness test 24 September 2020, 03:15:32 UTC
96f5a3c fix test for non-omnistaging 24 September 2020, 02:39:22 UTC
c42d736 remove limit on size of random arrays 24 September 2020, 02:37:34 UTC
0d5f15f Fix the abstract eval and translation rule for all_to_all The previous rules assumed that `split_axis == concat_axis` (i.e. that the used collective is equivalent to `pswapaxes`). Since we expose this as part of our API, we should probably make sure that we handle other cases too. Fixes #1332. 23 September 2020, 16:22:03 UTC
c875ab3 Merge pull request #4391 from apaszke:axis_index_handle_list PiperOrigin-RevId: 333304709 23 September 2020, 16:11:33 UTC
c9b3df3 Merge pull request #4382 from hawkinsp:lu PiperOrigin-RevId: 333296510 23 September 2020, 15:26:03 UTC
c6cd2f9 Merge pull request #4392 from SIben:fix_eigh_flakiness_cpu PiperOrigin-RevId: 333294460 23 September 2020, 15:13:34 UTC
back to top