| 38ce6d0 | Yash Katariya | 04 May 2022, 21:39:00 UTC | Update TF commit for release PiperOrigin-RevId: 446555288 | 04 May 2022, 21:42:50 UTC |
| 7297115 | jax authors | 04 May 2022, 21:37:25 UTC | Merge pull request #10546 from jakevdp:unravel-indices PiperOrigin-RevId: 446553390 | 04 May 2022, 21:37:25 UTC |
| 3c2d2b2 | Jake VanderPlas | 04 May 2022, 20:32:16 UTC | jnp.unravel_index: improve test coverage | 04 May 2022, 20:32:16 UTC |
| 58320e2 | Jake VanderPlas | 03 May 2022, 19:12:06 UTC | jnp.unravel_index: avoid overflow for large dimension sizes | 04 May 2022, 20:12:29 UTC |
| a8c6742 | Peter Hawkins | 04 May 2022, 19:33:40 UTC | Restrict Bazel visibility of //jaxlib:gpu_support PiperOrigin-RevId: 446524414 | 04 May 2022, 19:34:08 UTC |
| 03c8020 | jax authors | 04 May 2022, 19:09:51 UTC | Merge pull request #10532 from mattjj:remove-units-final PiperOrigin-RevId: 446519086 | 04 May 2022, 19:09:51 UTC |
| 8c6f916 | Peter Hawkins | 04 May 2022, 18:18:28 UTC | Reapply https://github.com/google/jax/pull/10482 now the TF PR is ready for submission. Make a couple of style cleanups in passing. PiperOrigin-RevId: 446506301 | 04 May 2022, 18:18:58 UTC |
| 9cd55a2 | Matthew Johnson | 03 May 2022, 00:11:44 UTC | [remove-units] remove units | 04 May 2022, 17:58:56 UTC |
| 97b7fd7 | jax authors | 04 May 2022, 17:51:01 UTC | Merge pull request #10564 from mattjj:gate-xla-version PiperOrigin-RevId: 446498170 | 04 May 2022, 17:51:01 UTC |
| 0c5864a | Matthew Johnson | 04 May 2022, 16:42:20 UTC | add xla_client._version checks for mhlo.ConstOp signature fix break from 0cf08d0c6841332240cae873e4b4cf9a9b313373 | 04 May 2022, 16:54:06 UTC |
| 0cf08d0 | jax authors | 04 May 2022, 12:31:05 UTC | Integrate LLVM at llvm/llvm-project@46cc04de341b Updates LLVM usage to match [46cc04de341b](https://github.com/llvm/llvm-project/commit/46cc04de341b) PiperOrigin-RevId: 446430294 | 04 May 2022, 12:31:41 UTC |
| 5c838d2 | Jean-Baptiste Lespiau | 04 May 2022, 08:21:39 UTC | Add an option when lowering to not remove unused arguments. This way, code using the output xla executable does not need to also drop the unused arguments, simplifying downstream code. PiperOrigin-RevId: 446391558 | 04 May 2022, 08:22:14 UTC |
| 35c0862 | Xin Zhou | 04 May 2022, 04:57:56 UTC | [mhlo] Add result type inference for mhlo.dynamic-slice. PiperOrigin-RevId: 446366018 | 04 May 2022, 04:58:29 UTC |
| 2498af3 | jax authors | 04 May 2022, 04:17:55 UTC | Merge pull request #10557 from sharadmv:dont-leak PiperOrigin-RevId: 446361367 | 04 May 2022, 04:17:55 UTC |
| 78a4e30 | Sharad Vikram | 04 May 2022, 03:53:30 UTC | Don't leak the keepalive in debug_callback lowering | 04 May 2022, 03:53:30 UTC |
| a9c0a97 | jax authors | 04 May 2022, 03:45:54 UTC | Merge pull request #10553 from sharadmv:debug-print PiperOrigin-RevId: 446357443 | 04 May 2022, 03:45:54 UTC |
| 2bb7e97 | jax authors | 04 May 2022, 03:29:22 UTC | Merge pull request #10549 from sharadmv:attach-executable PiperOrigin-RevId: 446355430 | 04 May 2022, 03:29:22 UTC |
| 7a641de | Sharad Vikram | 03 May 2022, 23:50:11 UTC | Add initial debug print implementation Co-authored-by: Matthew Johnson <mattjj@google.com> | 04 May 2022, 03:16:05 UTC |
| cfe9256 | Yash Katariya | 04 May 2022, 03:07:52 UTC | Fix the release bug where jaxlib version was not right. PiperOrigin-RevId: 446352145 | 04 May 2022, 03:08:24 UTC |
| ef982cf | Sharad Vikram | 14 April 2022, 21:18:31 UTC | Attach keepalive to executable | 04 May 2022, 03:08:03 UTC |
| c823025 | Yash Katariya | 04 May 2022, 02:03:59 UTC | Fix the checking of all lockfiles written. Previously the check was dependent on `gfile.listdir(lockfiles_dir)` in process 0. But other processes could run ahead and delete their lockfiles once done causing the check on process 0 to fail. Now process 0 keeps a count after a successful write and then checks against that count. PiperOrigin-RevId: 446343403 | 04 May 2022, 02:04:39 UTC |
| 3bfcabb | jax authors | 03 May 2022, 23:30:41 UTC | Merge pull request #10295 from sharadmv:jaxpr-effect-lowering PiperOrigin-RevId: 446317877 | 03 May 2022, 23:30:41 UTC |
| 8031eee | Sharad Vikram | 14 April 2022, 21:18:31 UTC | Add in runtime tokens for effectful jaxprs | 03 May 2022, 22:55:07 UTC |
| ff1a3c4 | Yash Katariya | 03 May 2022, 21:52:12 UTC | jax and jaxlib release PiperOrigin-RevId: 446295827 | 03 May 2022, 21:52:40 UTC |
| c6343dd | Jake VanderPlas | 03 May 2022, 20:47:11 UTC | jax.scipy.linalg.schur: error on 16-bit floats Fixes https://github.com/google/jax/issues/10530 PiperOrigin-RevId: 446279906 | 03 May 2022, 20:47:44 UTC |
| 37ea024 | Peter Hawkins | 03 May 2022, 18:46:10 UTC | Temporarily revert: 5d742fc0c3313798ba01b45bc13f973ef950e2e7 by lipracer <lipracer@gmail.com>: Compatible with RngBitGeneratorOp builder modifications I will reapply this PR when the TF PR is merged. PiperOrigin-RevId: 446250797 | 03 May 2022, 18:46:44 UTC |
| 02b9909 | Jean-Baptiste Lespiau | 03 May 2022, 07:56:38 UTC | Make `_one_replica_buffer_indices` callable by itself. PiperOrigin-RevId: 446131175 | 03 May 2022, 07:57:13 UTC |
| f7296a4 | Tianjian Lu | 03 May 2022, 06:52:31 UTC | [linalg] Update LU decomposition. PiperOrigin-RevId: 446122667 | 03 May 2022, 06:53:03 UTC |
| 31be621 | Yash Katariya | 03 May 2022, 06:38:45 UTC | Update the serialization example PiperOrigin-RevId: 446121048 | 03 May 2022, 06:39:21 UTC |
| c721544 | jax authors | 03 May 2022, 06:24:19 UTC | better error message in mesh_utils.py PiperOrigin-RevId: 446119369 | 03 May 2022, 06:24:56 UTC |
| de7a872 | Yash Katariya | 03 May 2022, 04:22:11 UTC | Take temp_checkpoint_dir and final_checkpoint_dir as the arguments to serialize instead of the __init__. THis is because this manager will be defined at the top where the directories may not yet be known. PiperOrigin-RevId: 446104174 | 03 May 2022, 04:22:43 UTC |
| 888e5c6 | Yash Katariya | 03 May 2022, 02:50:38 UTC | Update the version numbers after JAX release. PiperOrigin-RevId: 446092433 | 03 May 2022, 02:51:11 UTC |
| 634f58c | Peter Hawkins | 03 May 2022, 01:56:50 UTC | Enable a number of tests on GPU. In particular, pjit/xmap work on CPU these days. PiperOrigin-RevId: 446085110 | 03 May 2022, 01:57:27 UTC |
| b7293d5 | Yash Katariya | 03 May 2022, 01:43:27 UTC | Add fully asynchronous checkpointing. This will allow the training to proceed forward when the checkpoint is being committed. PiperOrigin-RevId: 446083057 | 03 May 2022, 01:43:54 UTC |
| 939233e | jax authors | 02 May 2022, 21:02:11 UTC | Merge pull request #10522 from mattjj:remove-units-partial-eval PiperOrigin-RevId: 446023568 | 02 May 2022, 21:02:11 UTC |
| 11ad045 | Matthew Johnson | 02 May 2022, 17:02:02 UTC | [remove-units] remove units from partial_eval.py After last week's changes, units are no longer traced or introduced into jaxprs in any way, so we don't need to use them in partial evaluation. (Also there are some unrelated removals of dead code in maps.py.) | 02 May 2022, 20:43:27 UTC |
| 44006c7 | jax authors | 02 May 2022, 20:37:30 UTC | Merge pull request #10526 from hawkinsp:cudasm PiperOrigin-RevId: 446017399 | 02 May 2022, 20:37:30 UTC |
| 9fb9e12 | Peter Hawkins | 02 May 2022, 20:03:32 UTC | Don't include PTX for older GPU generations. See: https://github.com/tensorflow/tensorflow/pull/55613 For a CUDA build at head with the default compute capabilities, reduces wheel size from 141MB to 112MB. Don't redundantly specify default compute capabilities in .bazelrc and in build.py. | 02 May 2022, 20:27:37 UTC |
| 189bbfa | jax authors | 02 May 2022, 19:47:31 UTC | Merge pull request #10524 from jakevdp:update-jupytext PiperOrigin-RevId: 446005313 | 02 May 2022, 19:47:31 UTC |
| 3a2f87f | jax authors | 02 May 2022, 18:58:13 UTC | Merge pull request #10481 from jakevdp:bcoo-reshape PiperOrigin-RevId: 445992966 | 02 May 2022, 18:58:13 UTC |
| 1134ff5 | Jake VanderPlas | 02 May 2022, 18:22:31 UTC | CI: update jupytext to 0.13.8 | 02 May 2022, 18:22:31 UTC |
| be9aac1 | Peter Hawkins | 02 May 2022, 17:03:34 UTC | Relax test tolerance for flaky test. PiperOrigin-RevId: 445961557 | 02 May 2022, 17:08:17 UTC |
| 1a9a796 | Jake VanderPlas | 02 May 2022, 16:11:55 UTC | [sparse] implement sparse rule for lax.reshape_p | 02 May 2022, 16:11:55 UTC |
| 29a828a | jax authors | 30 April 2022, 17:41:37 UTC | Merge pull request #10488 from jakevdp:bcoo-repr PiperOrigin-RevId: 445651381 | 30 April 2022, 17:41:37 UTC |
| 768bac8 | jax authors | 30 April 2022, 04:18:40 UTC | Merge pull request #10512 from google:update-pypi PiperOrigin-RevId: 445573116 | 30 April 2022, 04:18:40 UTC |
| 838f225 | Matthew Johnson | 30 April 2022, 03:15:52 UTC | update version and changelog for pypi | 30 April 2022, 03:15:52 UTC |
| 290c90d | Matthew Johnson | 30 April 2022, 02:58:16 UTC | Trivial change for backward compatibility stub. PiperOrigin-RevId: 445564091 | 30 April 2022, 02:58:50 UTC |
| a682a37 | jax authors | 30 April 2022, 00:04:56 UTC | Merge pull request #10509 from mattjj:avoid-units-linear-transpose PiperOrigin-RevId: 445536008 | 30 April 2022, 00:04:56 UTC |
| c01fed9 | jax authors | 29 April 2022, 23:54:15 UTC | Merge pull request #10508 from mattjj:avoid-units-lax-reduce PiperOrigin-RevId: 445534295 | 29 April 2022, 23:54:15 UTC |
| e3a1b7f | jax authors | 29 April 2022, 23:49:19 UTC | Merge pull request #10507 from mattjj:avoid-units-new-remat PiperOrigin-RevId: 445534171 | 29 April 2022, 23:49:19 UTC |
| 874374c | Yash Katariya | 29 April 2022, 23:42:47 UTC | Raise a better error when assert fails in mesh_sharding_specs PiperOrigin-RevId: 445533883 | 29 April 2022, 23:49:05 UTC |
| 0208490 | Tianjian Lu | 29 April 2022, 23:42:08 UTC | [linalg] Add tpu svd lowering rule. PiperOrigin-RevId: 445533767 | 29 April 2022, 23:43:53 UTC |
| b90df4b | jax authors | 29 April 2022, 23:38:17 UTC | Merge pull request #10506 from mattjj:remove-partial-eval-to-jaxpr-dynamic PiperOrigin-RevId: 445532927 | 29 April 2022, 23:38:17 UTC |
| 65bff3c | Matthew Johnson | 29 April 2022, 23:36:57 UTC | [remove-units] avoid unit-generating function in jax.linear_transpose | 29 April 2022, 23:37:43 UTC |
| f970f58 | Matthew Johnson | 29 April 2022, 23:28:08 UTC | [remove-units] avoid unit-generating function in lax.reduce | 29 April 2022, 23:28:08 UTC |
| fb4731d | jax authors | 29 April 2022, 23:26:51 UTC | Merge pull request #10502 from jakevdp:fix-myst PiperOrigin-RevId: 445530774 | 29 April 2022, 23:26:51 UTC |
| 0bf3241 | Matthew Johnson | 29 April 2022, 22:50:25 UTC | [remove-units] remove now-dead flax helper function | 29 April 2022, 23:18:51 UTC |
| 58a2f47 | Matthew Johnson | 29 April 2022, 23:16:08 UTC | [remove-units] avoid units in new remat The new remat implementation called a unit-generating function in its transpose rule. But there's no need for that! | 29 April 2022, 23:16:08 UTC |
| d0799ac | jax authors | 29 April 2022, 23:09:37 UTC | Merge pull request #10505 from mattjj:remove-units-custom-jvp PiperOrigin-RevId: 445527649 | 29 April 2022, 23:09:37 UTC |
| 472f72a | Jake VanderPlas | 29 April 2022, 23:01:41 UTC | docs: pin myst-nb to 0.13.2 | 29 April 2022, 23:01:41 UTC |
| 85dcad3 | Matthew Johnson | 29 April 2022, 22:37:27 UTC | [remove-units] remove units from custom_jvp/vjp | 29 April 2022, 22:42:41 UTC |
| 227e525 | jax authors | 29 April 2022, 22:40:16 UTC | Merge pull request #10458 from carlosgmartin:random_orthogonal_unitary PiperOrigin-RevId: 445522278 | 29 April 2022, 22:40:16 UTC |
| 4d10a8c | jax authors | 29 April 2022, 22:28:31 UTC | Merge pull request #10501 from mattjj:remove-partial-eval-jaxpr PiperOrigin-RevId: 445519759 | 29 April 2022, 22:28:31 UTC |
| 8080883 | jax authors | 29 April 2022, 22:20:39 UTC | Merge pull request #10500 from jakevdp:result-type-weak PiperOrigin-RevId: 445518089 | 29 April 2022, 22:20:39 UTC |
| 5a3d2e3 | Matthew Johnson | 29 April 2022, 21:54:07 UTC | [remove-units] remove partial_eval_jaxpr (no callers!) | 29 April 2022, 21:54:07 UTC |
| 5c4636c | Peter Hawkins | 29 April 2022, 21:44:00 UTC | [JAX] Validate that platforms passed to MHLO lowering are known to exist. In at least one instance a user was passing a XLA client object rather than the name of a platform. PiperOrigin-RevId: 445510282 | 29 April 2022, 21:44:32 UTC |
| b0f1d61 | jax authors | 29 April 2022, 21:23:37 UTC | Merge pull request #10495 from jakevdp:qr-empty PiperOrigin-RevId: 445505576 | 29 April 2022, 21:23:37 UTC |
| 4717a6b | Kehang Han | 29 April 2022, 21:05:41 UTC | updates maps.mesh() context manager in jax_101/pjit documentation PiperOrigin-RevId: 445501580 | 29 April 2022, 21:06:11 UTC |
| b3c35a1 | jax authors | 29 April 2022, 20:05:17 UTC | Merge pull request #10497 from mattjj:remat-fix-4 PiperOrigin-RevId: 445486179 | 29 April 2022, 20:05:17 UTC |
| 1a46957 | Jake VanderPlas | 29 April 2022, 19:54:40 UTC | dtypes.result_type: add optional return_weak_type_flag argument | 29 April 2022, 19:54:40 UTC |
| ec8adb4 | Matthew Johnson | 29 April 2022, 19:42:19 UTC | skip tests | 29 April 2022, 19:42:19 UTC |
| dc6282c | Peter Hawkins | 29 April 2022, 19:06:56 UTC | Allow sharded_jit on CPU. sharded_jit won't exist for long, but it works on CPU these days so we may as well enable the test since it's fewer lines of code that way. PiperOrigin-RevId: 445475116 | 29 April 2022, 19:07:29 UTC |
| b276c31 | Carlos Martin | 29 April 2022, 18:20:50 UTC | Added random.orthogonal. | 29 April 2022, 18:20:50 UTC |
| 7e8dc74 | jax authors | 29 April 2022, 18:11:45 UTC | Merge pull request #10483 from mattjj:remove-units-xmap PiperOrigin-RevId: 445462243 | 29 April 2022, 18:11:45 UTC |
| b4a4ffd | Matthew Johnson | 29 April 2022, 18:03:15 UTC | disable for now | 29 April 2022, 18:03:15 UTC |
| 36b0506 | Matthew Johnson | 29 April 2022, 18:01:32 UTC | fix scan dce rule Co-authored-by: Roy Frostig <frostig@google.com> | 29 April 2022, 18:01:32 UTC |
| ec8252f | Matthew Johnson | 29 April 2022, 17:56:03 UTC | broken remat test! | 29 April 2022, 17:56:03 UTC |
| 17d0dd7 | jax authors | 29 April 2022, 17:16:57 UTC | Merge pull request #10489 from mattjj:remove-units-dropvar-aval PiperOrigin-RevId: 445449038 | 29 April 2022, 17:16:57 UTC |
| f6dca14 | Jake VanderPlas | 29 April 2022, 17:04:47 UTC | Make lax.linalg.qr robust to zero-dimensional inputs | 29 April 2022, 17:04:47 UTC |
| 1788508 | Matthew Johnson | 29 April 2022, 16:06:02 UTC | reviewer comments | 29 April 2022, 16:09:07 UTC |
| b81f57b | jax authors | 29 April 2022, 15:37:08 UTC | Merge pull request #10482 from lipracer:main PiperOrigin-RevId: 445427001 | 29 April 2022, 15:37:08 UTC |
| 477dfa6 | Matthew Johnson | 29 April 2022, 05:51:41 UTC | [remove-units] don't use abstract_unit for dropvar avals | 29 April 2022, 05:51:41 UTC |
| 0bb37e2 | Jake VanderPlas | 28 April 2022, 23:52:17 UTC | [sparse] display n_batch & n_dense in BCOO repr | 28 April 2022, 23:52:17 UTC |
| cdd1167 | Tianjian Lu | 28 April 2022, 21:46:08 UTC | [signal] Update signal detrend test. PiperOrigin-RevId: 445253797 | 28 April 2022, 21:47:07 UTC |
| 4d41111 | Tianjian Lu | 28 April 2022, 21:41:29 UTC | [linalg] Update svd test on reconstructed operands and unitary singular vectors. PiperOrigin-RevId: 445252682 | 28 April 2022, 21:42:02 UTC |
| 4dad41f | jax authors | 28 April 2022, 20:34:46 UTC | Merge pull request #10484 from mattjj:pmap-redundant-axis-env-extension PiperOrigin-RevId: 445235173 | 28 April 2022, 20:34:46 UTC |
| a9f64a9 | Peter Hawkins | 28 April 2022, 20:16:14 UTC | Drop Bazel visibility of //jaxlib. PiperOrigin-RevId: 445230768 | 28 April 2022, 20:16:54 UTC |
| ca112da | Matthew Johnson | 28 April 2022, 05:04:34 UTC | [remove-units] avoid making xmap partial eval deal with units | 28 April 2022, 19:48:14 UTC |
| 8915391 | Matthew Johnson | 28 April 2022, 19:44:19 UTC | fix redundant (harmless) axis env extension in pmap partial eval | 28 April 2022, 19:46:19 UTC |
| 5d742fc | lipracer | 28 April 2022, 14:22:17 UTC | Compatible with RngBitGeneratorOp builder modifications | 28 April 2022, 19:10:39 UTC |
| e5f777e | jax authors | 28 April 2022, 18:22:37 UTC | Merge pull request #10466 from jakevdp:bcoo-concatenate PiperOrigin-RevId: 445202634 | 28 April 2022, 18:22:37 UTC |
| 2d9af38 | Jake VanderPlas | 28 April 2022, 17:59:11 UTC | [sparse] implement sparse rule for lax.concatenate_p | 28 April 2022, 17:59:11 UTC |
| 2864d58 | Peter Hawkins | 28 April 2022, 17:45:35 UTC | Lock down the default Bazel visibility of //jaxlib targets. PiperOrigin-RevId: 445192548 | 28 April 2022, 17:46:15 UTC |
| a9ae7ca | Peter Hawkins | 28 April 2022, 17:24:30 UTC | Reexport jaxlib.__version as jax.lib.__version__. PiperOrigin-RevId: 445186919 | 28 April 2022, 17:25:06 UTC |
| 0b47036 | Peter Hawkins | 28 April 2022, 13:01:22 UTC | Change the default jnp.take mode to "fill". Previously, `jnp.take` defaulted to clamping out-of-bounds indices into range. Now, `jnp.take` returns invalid values (e.g., NaN) for out-of-bounds indices. This change attempts to prevent latent bugs caused by inadvertent out-of-bounds indices. The previous behavior can be approximated using the "clip" or "wrap" fill modes. PiperOrigin-RevId: 445130143 | 28 April 2022, 13:01:56 UTC |
| 611759d | jax authors | 28 April 2022, 05:27:06 UTC | Merge pull request #8554 from mattjj:scan-dce-rule PiperOrigin-RevId: 445059933 | 28 April 2022, 05:27:06 UTC |
| 4608d36 | Matthew Johnson | 16 November 2021, 06:36:39 UTC | add scan dce rule | 28 April 2022, 03:47:43 UTC |
| b852778 | jax authors | 28 April 2022, 03:27:33 UTC | Merge pull request #10469 from mattjj:remove-units-pjit PiperOrigin-RevId: 445041987 | 28 April 2022, 03:27:33 UTC |
| 7a5a688 | Matthew Johnson | 28 April 2022, 02:58:33 UTC | [remove-units] prevent pjit partial eval from dealing with units | 28 April 2022, 03:07:20 UTC |
| a161d6a | jax authors | 28 April 2022, 02:49:44 UTC | Merge pull request #10468 from mattjj:remove-units-remat PiperOrigin-RevId: 445035270 | 28 April 2022, 02:49:44 UTC |
