| 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 |
| c7f2712 | jax authors | 21 September 2022, 16:47:21 UTC | Flip default value of jax_unique_mhlo_module_names to False. This should help avoid unnecessary cache misses. PiperOrigin-RevId: 475852954 | 21 September 2022, 16:48:01 UTC |
| d0e1c3e | Peter Hawkins | 21 September 2022, 15:50:17 UTC | Disable tests under sanitizers that are timing out in CI. PiperOrigin-RevId: 475839926 | 21 September 2022, 15:50:55 UTC |
| fd90f40 | jax authors | 21 September 2022, 13:12:44 UTC | Merge pull request #12443 from cloudhan:fix-mlir-chlo-stablehlo-symbols PiperOrigin-RevId: 475808753 | 21 September 2022, 13:12:44 UTC |
| dc414c8 | Cloud Han | 21 September 2022, 09:08:13 UTC | Switch from experimental_strict_action_env to incompatible_strict_action_env to avoid deprecation warning | 21 September 2022, 09:26:21 UTC |
| 3fa2c93 | Cloud Han | 21 September 2022, 09:07:15 UTC | Fix linker error due to chlo and stablehol symbols are not exported in mlir dll | 21 September 2022, 09:26:21 UTC |
| 6183727 | Yash Katariya | 20 September 2022, 23:38:09 UTC | Update pjit_test to skip GDA tests with Array is enabled. PiperOrigin-RevId: 475684445 | 20 September 2022, 23:38:43 UTC |
| 310bcd5 | jax authors | 20 September 2022, 18:23:42 UTC | Merge pull request #12389 from LenaMartens:check-while-2 PiperOrigin-RevId: 475606527 | 20 September 2022, 18:23:42 UTC |
| 018e700 | lenamartens | 16 September 2022, 16:46:25 UTC | Checkify: support batched while. | 20 September 2022, 16:59:46 UTC |
| 60c5b32 | jax authors | 20 September 2022, 16:59:22 UTC | Merge pull request #12418 from jakevdp:jep-update PiperOrigin-RevId: 475582299 | 20 September 2022, 16:59:22 UTC |
| e855a9c | jax authors | 20 September 2022, 16:52:56 UTC | Merge pull request #12428 from jakevdp:tracer-methods PiperOrigin-RevId: 475580916 | 20 September 2022, 16:52:56 UTC |
| fce1099 | Jake VanderPlas | 20 September 2022, 16:44:29 UTC | Update JEP-12049 implementation discussion | 20 September 2022, 16:44:29 UTC |
| 7469804 | Jake VanderPlas | 20 September 2022, 16:09:23 UTC | Tracer: add missing __round__ and __reversed__ methods | 20 September 2022, 16:09:23 UTC |
| fc2902c | Yash Katariya | 20 September 2022, 15:24:08 UTC | Make the gda and xmap sharding check work generally by checking the OpSharding protos. PiperOrigin-RevId: 475560097 | 20 September 2022, 15:24:47 UTC |
| 24bc153 | jax authors | 20 September 2022, 03:13:21 UTC | Merge pull request #12425 from sharadmv:vis-pmap PiperOrigin-RevId: 475449377 | 20 September 2022, 03:13:21 UTC |
| 0276a6e | Sharad Vikram | 20 September 2022, 02:19:45 UTC | Add support for pmap sharding | 20 September 2022, 02:29:44 UTC |
| 09d3ee1 | jax authors | 20 September 2022, 02:27:19 UTC | Merge pull request #12424 from sharadmv:vis-fixes PiperOrigin-RevId: 475442910 | 20 September 2022, 02:27:19 UTC |
| f825a3c | Sharad Vikram | 20 September 2022, 01:35:22 UTC | Limit console width for visualize_sharding | 20 September 2022, 01:41:45 UTC |
| e41e8d9 | Yash Katariya | 19 September 2022, 23:58:46 UTC | Only copy_to_device if the indices match. Otherwise reshard the array if its uncommitted. This is important where you have 1 process per device. PiperOrigin-RevId: 475418561 | 19 September 2022, 23:59:14 UTC |
| 441f400 | jax authors | 19 September 2022, 21:36:21 UTC | Merge pull request #12386 from sharadmv:viz_sharding PiperOrigin-RevId: 475387460 | 19 September 2022, 21:36:21 UTC |
