| 0282b4b | jax authors | 28 September 2022, 15:00:20 UTC | Merge pull request #12538 from jakevdp:bundle-pyi PiperOrigin-RevId: 477453094 | 28 September 2022, 15:00:20 UTC |
| aafc70d | jax authors | 28 September 2022, 13:50:19 UTC | Merge pull request #12556 from hawkinsp:rocm PiperOrigin-RevId: 477440001 | 28 September 2022, 13:50:19 UTC |
| 5fe7a54 | jax authors | 28 September 2022, 13:50:06 UTC | Merge pull request #12555 from hawkinsp:release PiperOrigin-RevId: 477439236 | 28 September 2022, 13:50:06 UTC |
| 39eabe8 | jax authors | 28 September 2022, 13:43:45 UTC | Merge pull request #12552 from hawkinsp:nccl PiperOrigin-RevId: 477439228 | 28 September 2022, 13:43:45 UTC |
| f7bafb3 | Peter Hawkins | 28 September 2022, 13:40:57 UTC | Disable multiprocess_gpu_test that fails on ROCm. | 28 September 2022, 13:40:57 UTC |
| 8d86436 | Peter Hawkins | 28 September 2022, 13:33:52 UTC | jax/jaxlib 0.3.20 release candidate. | 28 September 2022, 13:33:52 UTC |
| eabb91e | Peter Hawkins | 28 September 2022, 13:06:04 UTC | Fix test failure in GPU CI if NCCL_DEBUG is enabled. If NCCL_DEBUG is enabled, NCCL prints extra status information. Make test accept this. | 28 September 2022, 13:06:04 UTC |
| 96abd9a | jax authors | 28 September 2022, 05:33:12 UTC | Merge pull request #12540 from sharadmv:cond-lowering-fix PiperOrigin-RevId: 477358889 | 28 September 2022, 05:33:12 UTC |
| 96a85bd | Yash Katariya | 28 September 2022, 05:26:51 UTC | Make addressable_shards a property like local_shards PiperOrigin-RevId: 477358276 | 28 September 2022, 05:27:19 UTC |
| 9489068 | jax authors | 28 September 2022, 05:16:02 UTC | Merge pull request #12546 from mattjj:issue12542 PiperOrigin-RevId: 477356925 | 28 September 2022, 05:16:02 UTC |
| ddeaa8d | Sharad Vikram | 27 September 2022, 22:19:03 UTC | Fix lowering bug in effectful batched cond and add tests | 28 September 2022, 05:12:13 UTC |
| b4e1d0a | Yash Katariya | 28 September 2022, 04:31:58 UTC | Propagate `name` through ExecuteReplicated for `dispatch.check_special` PiperOrigin-RevId: 477351323 | 28 September 2022, 04:32:32 UTC |
| b175e11 | Matthew Johnson | 28 September 2022, 03:39:19 UTC | [c++ jit] only set use_fastpath in cache_miss if all args are DeviceArrays fixes #12542 Co-authored-by: Peter Hawkins <phawkins@google.com> Co-authored-by: Kuangyuan Chen <chky@google.com> | 28 September 2022, 03:51:07 UTC |
| 933b6a2 | Yash Katariya | 28 September 2022, 03:27:11 UTC | Fix the bug where XLA doesn't provide shardings for all the outputs if all the elements in the output tuple have the same sharding. XLA decides to run the `FusionTupleDeduplicator` to put the sharding on ROOT instead of the tuple. PiperOrigin-RevId: 477343328 | 28 September 2022, 03:27:39 UTC |
| c8bff11 | Yash Katariya | 28 September 2022, 02:18:49 UTC | Add `addressable_` counterparts of `local_` to GDA to make it easier for users to move to Array as both will have the same API. PiperOrigin-RevId: 477332697 | 28 September 2022, 02:19:29 UTC |
| e4f2bff | Yash Katariya | 27 September 2022, 23:01:46 UTC | Disintegrate `Array` into DeviceBuffers inside GDA. This is required for backwards compatibility changes as users can create GDAs and pass that to pjit even when Array is switched on. PiperOrigin-RevId: 477297406 | 27 September 2022, 23:02:23 UTC |
| 0919a67 | jax authors | 27 September 2022, 20:31:05 UTC | Merge pull request #12534 from google:update-pypi PiperOrigin-RevId: 477260550 | 27 September 2022, 20:31:05 UTC |
| 6e6fb10 | Jake VanderPlas | 27 September 2022, 19:55:20 UTC | setup: bundle *.pyi files with distribution | 27 September 2022, 19:55:42 UTC |
| d028d93 | Skye Wanderman-Milne | 27 September 2022, 18:00:27 UTC | Update version and changelog for jax 0.3.19 release | 27 September 2022, 18:00:27 UTC |
| 9e4114f | Yash Katariya | 27 September 2022, 17:06:10 UTC | Move `array.py` and `sharding.py` from `experimental/` to `_src/`. PiperOrigin-RevId: 477201711 | 27 September 2022, 17:06:52 UTC |
| 0e11688 | jax authors | 27 September 2022, 15:38:46 UTC | Merge pull request #12382 from jakevdp:reduction-dtype PiperOrigin-RevId: 477179725 | 27 September 2022, 15:38:46 UTC |
| 1bcf8d6 | jax authors | 27 September 2022, 01:14:56 UTC | Merge pull request #12497 from mattjj:djax-dag-fix1 PiperOrigin-RevId: 477038279 | 27 September 2022, 01:14:56 UTC |
| e42247b | jax authors | 27 September 2022, 01:08:45 UTC | Merge pull request #12524 from sharadmv:lax-import-fix PiperOrigin-RevId: 477038211 | 27 September 2022, 01:08:45 UTC |
| 389a2e5 | Yash Katariya | 27 September 2022, 00:49:48 UTC | Add a backwards compat path for `op_sharding.clone()` because it doesn't exist with the latest jaxlib on pypi PiperOrigin-RevId: 477034758 | 27 September 2022, 00:50:19 UTC |
| 1e7ca8f | Matthew Johnson | 23 September 2022, 21:21:18 UTC | fix bug in djax type signature inference logic Co-authored-by: Sharad Vikram <sharad.vikram@gmail.com> | 27 September 2022, 00:48:25 UTC |
| 1d895b2 | Sharad Vikram | 27 September 2022, 00:29:08 UTC | Fix lax imports | 27 September 2022, 00:32:44 UTC |
| cbf34cb | Yash Katariya | 26 September 2022, 23:17:26 UTC | Rename the concrete class `Array` to `ArrayImpl` PiperOrigin-RevId: 477017236 | 26 September 2022, 23:18:30 UTC |
| 71bcabe | Tianjian Lu | 26 September 2022, 23:01:47 UTC | [sparse] Add BCSR format template. PiperOrigin-RevId: 477013899 | 26 September 2022, 23:02:16 UTC |
| 82636b0 | jax authors | 26 September 2022, 22:23:02 UTC | Merge pull request #12523 from jakevdp:fix-build PiperOrigin-RevId: 477005157 | 26 September 2022, 22:23:02 UTC |
| 6cae54f | Jake VanderPlas | 26 September 2022, 22:13:12 UTC | Fix bazel build alias | 26 September 2022, 22:13:12 UTC |
| d63a944 | Peter Hawkins | 26 September 2022, 21:38:06 UTC | Change jax_jit_test to be a jax_test() under Bazel that works across backends. Make it pass under TPU if x64 types are enabled. PiperOrigin-RevId: 476994286 | 26 September 2022, 21:38:35 UTC |
| 265b39d | Jake VanderPlas | 26 September 2022, 21:17:43 UTC | Add pytype_srcs to main jax BUILD rule PiperOrigin-RevId: 476989241 | 26 September 2022, 21:18:13 UTC |
| ddd8581 | jax authors | 26 September 2022, 20:41:31 UTC | Merge pull request #12480 from google:bug-template-gpu-smi PiperOrigin-RevId: 476979981 | 26 September 2022, 20:41:31 UTC |
| 1860f6d | Jake VanderPlas | 26 September 2022, 20:31:43 UTC | [x64] add promote_integers argument to jnp.prod & jnp.sum | 26 September 2022, 20:31:43 UTC |
| 69d1a2c | jax authors | 26 September 2022, 20:00:31 UTC | Merge pull request #12517 from skye:update-pypi PiperOrigin-RevId: 476969287 | 26 September 2022, 20:00:31 UTC |
| b2b60d9 | Yash Katariya | 26 September 2022, 19:43:13 UTC | Add `make_array_from_single_device_arrays` to prepare to rename of the concrete `Array` to `ArrayImpl`. PiperOrigin-RevId: 476965287 | 26 September 2022, 19:43:59 UTC |
| 3c0d280 | Skye Wanderman-Milne | 26 September 2022, 19:38:32 UTC | Update version and changelog for jax 0.3.18 release | 26 September 2022, 19:43:39 UTC |
| 2a7b319 | Roy Frostig | 26 September 2022, 18:06:29 UTC | add `nvidia-smi` question to bug template | 26 September 2022, 18:06:29 UTC |
| e034432 | jax authors | 26 September 2022, 17:04:14 UTC | Merge pull request #12513 from inoryy:patch-4 PiperOrigin-RevId: 476923412 | 26 September 2022, 17:04:14 UTC |
| 7962b01 | jax authors | 26 September 2022, 16:53:40 UTC | Merge pull request #12485 from LenaMartens:checkify-lower PiperOrigin-RevId: 476922387 | 26 September 2022, 16:53:40 UTC |
| 27e3981 | lenamartens | 23 September 2022, 14:10:41 UTC | lowerable errors behind a config flag. | 26 September 2022, 16:34:27 UTC |
| 8bcf358 | Roman Ring | 26 September 2022, 16:14:09 UTC | Remove unused _remat_static_argnums import. | 26 September 2022, 16:14:09 UTC |
| 78ecc14 | lenamartens | 23 September 2022, 13:30:49 UTC | Lowerable checks!! | 26 September 2022, 15:54:18 UTC |
| 28672cc | jax authors | 26 September 2022, 15:50:13 UTC | Merge pull request #12496 from mattjj:improve-leak-checker-2 PiperOrigin-RevId: 476907407 | 26 September 2022, 15:50:13 UTC |
| 9c66569 | jax authors | 26 September 2022, 15:23:02 UTC | Merge pull request #12468 from LenaMartens:checkify-but-better PiperOrigin-RevId: 476901601 | 26 September 2022, 15:23:02 UTC |
| 2df61b1 | jax authors | 26 September 2022, 15:07:11 UTC | Merge pull request #12421 from jakevdp:jax-array PiperOrigin-RevId: 476898184 | 26 September 2022, 15:07:11 UTC |
| 0cb233e | Jake VanderPlas | 23 September 2022, 16:59:46 UTC | Add initial jax.Array base class for instance checks & annotation | 26 September 2022, 14:48:43 UTC |
| 2180710 | jax authors | 26 September 2022, 14:24:44 UTC | Merge pull request #12511 from hawkinsp:release PiperOrigin-RevId: 476889960 | 26 September 2022, 14:24:44 UTC |
| bcd36d8 | Peter Hawkins | 26 September 2022, 14:10:57 UTC | Jax and jaxlib 0.3.18 release candidate. | 26 September 2022, 14:10:57 UTC |
| 53de057 | jax authors | 26 September 2022, 13:58:46 UTC | Merge pull request #12510 from hawkinsp:context PiperOrigin-RevId: 476884674 | 26 September 2022, 13:58:46 UTC |
| f4bc663 | Peter Hawkins | 26 September 2022, 13:48:56 UTC | Wrap multiprocess test popen() uses in a context manager. Ensures that resources from popen() are cleaned up. | 26 September 2022, 13:48:56 UTC |
| ec15e83 | jax authors | 26 September 2022, 03:53:45 UTC | - Wraps calls to lax.xeinsum and _einsum in a named call with their 'spec', the string specifying the computation. Makes xprof traces more interpretable. PiperOrigin-RevId: 476796185 | 26 September 2022, 03:54:17 UTC |
| 7c85ca3 | Yash Katariya | 24 September 2022, 00:31:00 UTC | Only look at hlo_modules for output sharding if there is more than 1 device because if there is only 1 device, the spmd partitioner won't run. PiperOrigin-RevId: 476497929 | 24 September 2022, 00:31:33 UTC |
| 8ee7129 | Peter Hawkins | 24 September 2022, 00:11:10 UTC | Fix jnp.unwrap() test failures on GPU. A recent XLA change allows XLA to use excess precision on GPU, which caused CompileAndCheck to report noticeable numerical changes for bfloat16. In passing, also enable the comparison against NumPy test for bfloat16 by using a wrapper function. PiperOrigin-RevId: 476494989 | 24 September 2022, 00:11:51 UTC |
| d2fcfb6 | jax authors | 23 September 2022, 21:51:11 UTC | Merge pull request #12407 from hirwa-nshuti:docs-fix PiperOrigin-RevId: 476467728 | 23 September 2022, 21:51:11 UTC |
| 03abcc7 | Matthew Johnson | 23 September 2022, 19:42:15 UTC | fix typo in test | 23 September 2022, 21:43:24 UTC |
| e76aa77 | jax authors | 23 September 2022, 20:38:59 UTC | Merge pull request #12437 from sudhakarsingh27:add_multi_host_pjit_tests PiperOrigin-RevId: 476451469 | 23 September 2022, 20:38:59 UTC |
| 1fa0dda | Yash Katariya | 23 September 2022, 20:29:47 UTC | Return single device Arrays from `.device_buffer` and `.device_buffers`. PiperOrigin-RevId: 476449591 | 23 September 2022, 20:30:26 UTC |
| 43bbce0 | jax authors | 23 September 2022, 20:09:26 UTC | Merge pull request #12486 from hawkinsp:debugging PiperOrigin-RevId: 476445041 | 23 September 2022, 20:09:26 UTC |
| 737327a | jax authors | 23 September 2022, 19:58:03 UTC | Merge pull request #12490 from mattjj:improve-leak-checker PiperOrigin-RevId: 476442352 | 23 September 2022, 19:58:03 UTC |
| b6ef90f | Matthew Johnson | 23 September 2022, 18:24:13 UTC | fix leak checker internal error The issue was that partial_eval.py's _memoize, used in custom_jvp, was made into an identity function by enabling config.jax_check_tracer_leaks (from references to the main trace (needed for the jvp_jaxpr thunk) and hence trigger the leak checker (which would see if any references to the main trace persisted after finishing tracing of the user function). But after #7345, the leak checker should only trigger when actual Tracers are leaked. So disabling the memoization when jax_check_tracer_leaks is no longer active shouldn't be necessary. (These PR numbers seem out of order! We're not sure why.) Co-authored-by: Sharad Vikram <sharad.vikram@gmail.com> | 23 September 2022, 19:33:45 UTC |
| 4dd0d85 | Sudhakar | 23 September 2022, 19:11:56 UTC | add multihost pjit tests | 23 September 2022, 19:11:56 UTC |
| a6b24b3 | Jake VanderPlas | 23 September 2022, 19:06:35 UTC | Add regression test for lax.rev simplification error PiperOrigin-RevId: 476430486 | 23 September 2022, 19:07:15 UTC |
| ecb27a9 | Yash Katariya | 23 September 2022, 18:40:01 UTC | Update the `_check_special` code to not use xla_shape since its deprecated and does not work with Array. PiperOrigin-RevId: 476422732 | 23 September 2022, 18:40:32 UTC |
| d078f3f | jax authors | 23 September 2022, 18:31:37 UTC | Merge pull request #12478 from sharadmv:sharding-docs PiperOrigin-RevId: 476420315 | 23 September 2022, 18:31:37 UTC |
| e8865c8 | jax authors | 23 September 2022, 17:55:10 UTC | Merge pull request #12481 from kho:changelist/476272494 PiperOrigin-RevId: 476411483 | 23 September 2022, 17:55:10 UTC |
| c823151 | Ke Wu | 23 September 2022, 04:33:07 UTC | Allow transpose axes to be negative to match (undocumented) NumPy behavior | 23 September 2022, 17:18:23 UTC |
| 0c08547 | Tres Popp | 23 September 2022, 16:34:12 UTC | Modify CorrCoef test to not rely on floating poing representation of 1/3 The operation computed an average while using the dimension of size 3. This is then changed to multiplying by 1/3 with compilers, but 1/3 cannot be represented perfectly. That made this test case rely on a very precise result from an unrepresentable calculation. PiperOrigin-RevId: 476391389 | 23 September 2022, 16:39:01 UTC |
| 38fb8ed | Peter Hawkins | 23 September 2022, 16:32:00 UTC | Fix copyright attribution for some newly added files. PiperOrigin-RevId: 476390902 | 23 September 2022, 16:32:47 UTC |
| 6c47dc5 | jax authors | 23 September 2022, 16:16:38 UTC | Merge pull request #12471 from ROCmSoftwarePlatform:rocm-dockerfile-update PiperOrigin-RevId: 476387200 | 23 September 2022, 16:16:38 UTC |
| da50bdd | Yash Katariya | 23 September 2022, 15:59:16 UTC | Fix the asan failure in pjit_test_cpu build target PiperOrigin-RevId: 476382929 | 23 September 2022, 15:59:57 UTC |
| c8f5541 | Yash Katariya | 23 September 2022, 15:47:45 UTC | Convert the devices in the `Mesh` constructor to a numpy array if its a list, tuple, etc. PiperOrigin-RevId: 476380496 | 23 September 2022, 15:48:31 UTC |
| a88c5ad | Peter Hawkins | 23 September 2022, 14:53:50 UTC | Fix xla extension version test in debugging.py The custom call partitioner callback was not present in version 94 but is present in version 95. | 23 September 2022, 14:53:50 UTC |
| 254dc24 | jax authors | 23 September 2022, 14:29:17 UTC | Merge pull request #11961 from jakeh-gc:plugin_device PiperOrigin-RevId: 476363760 | 23 September 2022, 14:29:17 UTC |
| 342f896 | jax authors | 23 September 2022, 14:13:50 UTC | Merge pull request #12484 from hawkinsp:doc PiperOrigin-RevId: 476361203 | 23 September 2022, 14:13:50 UTC |
| 7078f81 | lenamartens | 22 September 2022, 14:23:54 UTC | Checkify: misc improvements. - err.throw == check_error(err) -> meaning they have the same behavior under checkify now - "divided by zero" -> "division by zero" - add validation that check_error only takes args of type Error | 23 September 2022, 13:33:06 UTC |
| eed3279 | Peter Hawkins | 23 September 2022, 13:11:15 UTC | Improve documentation for unique_indices. | 23 September 2022, 13:11:15 UTC |
| 820efab | Felix Hirwa Nshuti | 23 September 2022, 06:23:09 UTC | removed repeated nan_to_num in docs | 23 September 2022, 06:23:09 UTC |
| 67b7ae2 | Tianjian Lu | 23 September 2022, 03:21:29 UTC | [sparse] Move `_bcoo_nse` to sparse util. PiperOrigin-RevId: 476263483 | 23 September 2022, 03:22:06 UTC |
| 99d4d8b | Sharad Vikram | 23 September 2022, 02:04:23 UTC | Update debugging docs to have sharding visualization | 23 September 2022, 02:42:36 UTC |
| 805073f | Sharad Vikram | 23 September 2022, 00:36:20 UTC | Add inspect_array_sharding, enabling looking at shardings in pjit-ted functions PiperOrigin-RevId: 476237731 | 23 September 2022, 00:36:56 UTC |
| 11a6fd9 | jax authors | 22 September 2022, 23:04:20 UTC | Merge pull request #12476 from jakevdp:match-sharding PiperOrigin-RevId: 476220059 | 22 September 2022, 23:04:20 UTC |
| 3d23592 | Jake VanderPlas | 22 September 2022, 22:28:59 UTC | [array] full_like: only match sharding if shape==None | 22 September 2022, 22:28:59 UTC |
| dfdf00c | jax authors | 22 September 2022, 21:00:07 UTC | Merge pull request #12472 from google:sharadmv-patch-2 PiperOrigin-RevId: 476190259 | 22 September 2022, 21:00:07 UTC |
| bc08381 | jax authors | 22 September 2022, 20:18:25 UTC | Merge pull request #12152 from nvcastet:add_slurm_orchestrator_support PiperOrigin-RevId: 476179963 | 22 September 2022, 20:18:25 UTC |
| 1a8a8a5 | Sharad Vikram | 22 September 2022, 19:55:55 UTC | Fix example in `pjit` docstring | 22 September 2022, 19:55:55 UTC |
| ba557d5 | Peter Hawkins | 22 September 2022, 19:26:48 UTC | Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.". See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details. PiperOrigin-RevId: 476167538 | 22 September 2022, 19:27:19 UTC |
| 412a537 | Nicolas Castet | 26 August 2022, 19:23:57 UTC | Add generic interface for auto initialization of distributed JAX service * Also add slurm cluster support | 22 September 2022, 19:15:38 UTC |
| d52de20 | Tyler Augustine | 22 September 2022, 18:44:22 UTC | Disable tests that timeout in debug mode in CI PiperOrigin-RevId: 476157051 | 22 September 2022, 18:44:56 UTC |
| 9a11b61 | Jason Furmanek | 22 September 2022, 18:29:30 UTC | [ROCM] Update Dockerfil.rocm to Ubuntu20 | 22 September 2022, 18:29:30 UTC |
| a157982 | Yash Katariya | 22 September 2022, 18:13:04 UTC | Make `jit(f).lower(*args)` go via lower_sharding_computation when `jax_array` is enabled. PiperOrigin-RevId: 476148608 | 22 September 2022, 18:13:33 UTC |
| 640e15f | Tres Popp | 22 September 2022, 08:28:45 UTC | Don't tuple arguments passed to XLA:CPU This is not needed and tuples are being avoided when possible for new code. This is tested by CPPJitTest.test_jit_with_many_args_works in jax/tests:api_test_cpu PiperOrigin-RevId: 476032228 | 22 September 2022, 08:29:14 UTC |
| 405a231 | Kuangyuan Chen | 22 September 2022, 03:17:38 UTC | Implement pjit fast path in cpp for jax.Array inputs PiperOrigin-RevId: 475988677 | 22 September 2022, 03:18:18 UTC |
| 52476d1 | Yash Katariya | 22 September 2022, 01:18:57 UTC | Add addressable_data to Array (similar to GDA) to aid in transition and also in auto spmd partitioner mode, always convert to MeshPspecSharding. PiperOrigin-RevId: 475972534 | 22 September 2022, 01:19:35 UTC |
| a09ef8a | Kuangyuan Chen | 22 September 2022, 01:05:23 UTC | Temporarily skip LaxBackedNumpyTests.testUnwrap on gpu to unblock jaxlib build PiperOrigin-RevId: 475970440 | 22 September 2022, 01:06:02 UTC |
| d41fa29 | jax authors | 21 September 2022, 20:26:01 UTC | Merge pull request #12370 from jakevdp:lax-sort-overload PiperOrigin-RevId: 475907384 | 21 September 2022, 20:26:01 UTC |
| 62d1783 | jax authors | 21 September 2022, 20:19:52 UTC | Merge pull request #12451 from jakevdp:array-declarations PiperOrigin-RevId: 475906903 | 21 September 2022, 20:19:52 UTC |
| 2dde633 | Jake VanderPlas | 21 September 2022, 19:51:32 UTC | [typing] add class-level declarations of Array members. This fixes some pytype errors associated with the changes in #12421 | 21 September 2022, 19:51:32 UTC |
| 541aadc | George Karpenkov | 21 September 2022, 19:04:03 UTC | [XLA:GPU] Allow simplifying lowering-precision-conversions by default This might lead to the output having higher precision than specified by HLO. PiperOrigin-RevId: 475889141 | 21 September 2022, 19:04:45 UTC |
| b9e3843 | Kuangyuan Chen | 21 September 2022, 16:48:32 UTC | Add many args benchmark for jax.Array PiperOrigin-RevId: 475853211 | 21 September 2022, 16:54:51 UTC |
