| 744f6b4 | Skye Wanderman-Milne | 16 May 2022, 19:07:15 UTC | Update xla_client._version and add missing version checks to JAX PiperOrigin-RevId: 449021408 | 16 May 2022, 19:07:47 UTC |
| bd20f0f | jax authors | 16 May 2022, 18:02:34 UTC | Merge pull request #10723 from hawkinsp:pm PiperOrigin-RevId: 449004204 | 16 May 2022, 18:02:34 UTC |
| a91fdf1 | jax authors | 16 May 2022, 17:31:12 UTC | Merge pull request #10700 from jakevdp:haskell-sig PiperOrigin-RevId: 448995683 | 16 May 2022, 17:31:12 UTC |
| 246f99c | jax authors | 16 May 2022, 17:21:10 UTC | Merge pull request #10693 from jakevdp:bcoo-update-method PiperOrigin-RevId: 448993136 | 16 May 2022, 17:21:10 UTC |
| d56601a | Peter Hawkins | 16 May 2022, 12:58:56 UTC | Include mlir.transforms and mlir.passmanager in jaxlib BUILD. These increase the binary size of jaxlib only a negligible amount, and allow running passes like canonicalization from Python. | 16 May 2022, 13:00:22 UTC |
| f73f03e | Christian Sigg | 16 May 2022, 07:39:10 UTC | Fix docs referring to `--print-ir-before/after-all`, which has changed to `--mlir-print...`. PiperOrigin-RevId: 448890307 | 16 May 2022, 07:39:40 UTC |
| 6a66052 | Yash Katariya | 16 May 2022, 01:35:15 UTC | Update values after jax release PiperOrigin-RevId: 448854487 | 16 May 2022, 01:35:46 UTC |
| f572646 | jax authors | 16 May 2022, 00:43:05 UTC | Merge pull request #10718 from hawkinsp:jaxfix PiperOrigin-RevId: 448849099 | 16 May 2022, 00:43:05 UTC |
| 337ec47 | Peter Hawkins | 16 May 2022, 00:24:04 UTC | Fix jax 0.3.11 GPU breakge when used with jaxlib 0.3.10. | 16 May 2022, 00:24:04 UTC |
| 1381afc | Yash Katariya | 15 May 2022, 19:13:54 UTC | Update version after jax release PiperOrigin-RevId: 448822949 | 15 May 2022, 19:14:26 UTC |
| 42421c8 | jax authors | 14 May 2022, 22:22:31 UTC | Merge pull request #10688 from sharadmv:color-by-default PiperOrigin-RevId: 448723842 | 14 May 2022, 22:22:31 UTC |
| b8a523f | Sharad Vikram | 13 May 2022, 05:22:52 UTC | Enable colors when we are using a terminal or IPython | 14 May 2022, 21:54:11 UTC |
| f26133c | jax authors | 14 May 2022, 21:37:46 UTC | Merge pull request #10584 from mattjj:remove-jnp-array-handling-raw-buffers PiperOrigin-RevId: 448720084 | 14 May 2022, 21:37:46 UTC |
| 86899ee | jax authors | 14 May 2022, 21:32:39 UTC | Merge pull request #10711 from mattjj:closed-call PiperOrigin-RevId: 448720052 | 14 May 2022, 21:32:39 UTC |
| 05dda56 | Matthew Johnson | 14 May 2022, 18:03:50 UTC | add core.closed_call_p | 14 May 2022, 21:06:30 UTC |
| c0d6a04 | Matthew Johnson | 05 May 2022, 05:59:06 UTC | remove jnp.array case for handling buffers w/ aval=None This functionality was added in #8134, but was superceded by later changes which ensured that we never produce DeviceArrays with their 'aval' property set to None (even when indexing ShardedDeviceArrays with integers, which used to be a problem case). | 14 May 2022, 15:21:54 UTC |
| ba0a2b3 | Hyojun Kim | 13 May 2022, 23:23:55 UTC | Add profile_version to CompileOptions profile_version is added to CompileOptions to support XLA compilation profile. PiperOrigin-RevId: 448599962 | 13 May 2022, 23:24:29 UTC |
| db73670 | Peter Hawkins | 13 May 2022, 20:56:57 UTC | Add support for padded arrays in QDWH algorithm. This change is in preparation for adding a jit-table QDWH-eig implementation. PiperOrigin-RevId: 448571523 | 13 May 2022, 20:57:36 UTC |
| cb26b72 | Jake VanderPlas | 06 May 2022, 23:46:05 UTC | [sparse] add update_layout() method | 13 May 2022, 19:39:21 UTC |
| 7c582ab | jax authors | 13 May 2022, 19:38:05 UTC | Merge pull request #10680 from jakevdp:rm-future PiperOrigin-RevId: 448555184 | 13 May 2022, 19:38:05 UTC |
| 8a3bfe0 | Jake VanderPlas | 13 May 2022, 19:35:28 UTC | DOC: add references for haskell-style signatures | 13 May 2022, 19:35:28 UTC |
| 4d22d22 | Jake VanderPlas | 13 May 2022, 18:58:10 UTC | Fix version skew in api_test.py Issue was introduced in https://github.com/google/jax/commit/8b073d482ec897e4cbda0e7650d6cfc9e8fb9811 PiperOrigin-RevId: 448546982 | 13 May 2022, 18:58:45 UTC |
| bd659df | jax authors | 13 May 2022, 17:19:23 UTC | Merge pull request #10686 from mattjj:tweak-mlir-shape-tensor-helper PiperOrigin-RevId: 448523702 | 13 May 2022, 17:19:23 UTC |
| 8b073d4 | Shuangchi He | 13 May 2022, 09:22:43 UTC | PR #55768: Fix typos for occured, appearence, this, is, a, for, agressiveness, t… Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/55768 Fix typos for occured, appearence, this, is, a, for, agressiveness, to, instrution, on. Copybara import of the project: -- 531b97d4b242a5642a221349ca0bd3132d6539a2 by Yulv-git <yulvchi@qq.com>: Fix typos for occured, appearence, this, is, a, for, agressiveness, to, instrution, on. Merging this change closes #55768 PiperOrigin-RevId: 448444384 | 13 May 2022, 09:23:14 UTC |
| d719e5a | Matthew Johnson | 12 May 2022, 23:55:50 UTC | tweak mlir shape_tensor helper, fewer MHLO ops | 13 May 2022, 00:02:08 UTC |
| 2100c3b | jax authors | 12 May 2022, 19:40:44 UTC | Merge pull request #10681 from JeppeKlitgaard:fix-int-typo PiperOrigin-RevId: 448313566 | 12 May 2022, 19:40:44 UTC |
| b9be65c | jax authors | 12 May 2022, 18:58:00 UTC | Merge pull request #10679 from jakevdp:fix-dollarmath PiperOrigin-RevId: 448303782 | 12 May 2022, 18:58:00 UTC |
| f92558c | Jake VanderPlas | 12 May 2022, 18:57:21 UTC | Remove unnecessary __future__ imports | 12 May 2022, 18:57:21 UTC |
| e422441 | Jeppe Klitgaard | 12 May 2022, 18:45:41 UTC | fix: mypy warning | 12 May 2022, 18:45:41 UTC |
| 6812402 | Jake VanderPlas | 12 May 2022, 17:37:34 UTC | DOC: explicitly enable dollarmath extension | 12 May 2022, 17:37:34 UTC |
| 41e1635 | jax authors | 12 May 2022, 17:30:11 UTC | Merge pull request #10664 from PhilipVinc:patch-1 PiperOrigin-RevId: 448282864 | 12 May 2022, 17:30:11 UTC |
| 0d88bdf | jax authors | 12 May 2022, 17:18:10 UTC | Merge pull request #10642 from lgeiger:reverse-trivial-gather PiperOrigin-RevId: 448279969 | 12 May 2022, 17:18:10 UTC |
| 7ba36fc | Peter Hawkins | 12 May 2022, 14:15:55 UTC | Change implementation of jax.scipy.linalg.polar() and jax._src.scipy.eigh to use the QDWH decomposition from jax._src.lax.qdwh. Remove jax._src.lax.polar. PiperOrigin-RevId: 448241206 | 12 May 2022, 14:20:52 UTC |
| 880cfc9 | Peter Hawkins | 12 May 2022, 14:15:17 UTC | Fix pseudocode in lax.switch() documentation to accept variadic operands. PiperOrigin-RevId: 448241093 | 12 May 2022, 14:15:46 UTC |
| 8e4cf25 | jax authors | 11 May 2022, 23:27:58 UTC | Merge pull request #10616 from mattjj:partial-eval-jaxpr-custom-upgrade PiperOrigin-RevId: 448110206 | 11 May 2022, 23:27:58 UTC |
| 6b2d1ba | jax authors | 11 May 2022, 23:23:06 UTC | Merge pull request #10595 from mattjj:core-test-with-corecall PiperOrigin-RevId: 448109351 | 11 May 2022, 23:23:06 UTC |
| 06e4fbc | Xin Zhou | 11 May 2022, 23:17:23 UTC | [mhlo] Add result type inference for mhlo.cholesky. PiperOrigin-RevId: 448108980 | 11 May 2022, 23:18:03 UTC |
| 2340725 | jax authors | 11 May 2022, 23:01:02 UTC | Merge pull request #10669 from mattjj:fix-negative-argnums PiperOrigin-RevId: 448105356 | 11 May 2022, 23:01:02 UTC |
| 0b841cf | Matthew Johnson | 05 May 2022, 22:18:57 UTC | make core_test.py pass with core.call | 11 May 2022, 22:45:40 UTC |
| f13b69c | Lukas Geiger | 09 May 2022, 22:58:42 UTC | Avoid generating trivial gathers when reversing array | 11 May 2022, 22:16:40 UTC |
| 5bce808 | jax authors | 11 May 2022, 21:52:43 UTC | Merge pull request #10671 from LionOfJewdah:retry_loop_fix PiperOrigin-RevId: 448090523 | 11 May 2022, 21:52:43 UTC |
| 4a63188 | David Silverstone | 11 May 2022, 21:19:28 UTC | Fixed a typo in `cloud_tpu_init.py` | 11 May 2022, 21:19:28 UTC |
| 7e241b6 | Matthew Johnson | 07 May 2022, 04:55:04 UTC | improve partial_eval_jaxpr_custom * add caching via weakref_lru_cache * add inst_in argument (needed for fixedpoints for loop primitives, in follow-up PR), update callers not to over-instantiate inputs (previously I had used a convention where call primitives would just stage out eqns with all inputs instantiated, for expediene) * add ensure_out_unknowns and ensure_out_inst arguments, analogues of `instantiate` on e.g. partial_eval_jaxpr, jvp_jaxpr, etc (also neede for fixpoints of loop primitives) * better dce in remat_partial_eval (e.g. prune unused residuals) | 11 May 2022, 20:20:23 UTC |
| 705e241 | Peter Hawkins | 11 May 2022, 20:06:19 UTC | Change non-array arguments to jax.lax.linalg functions to be keyword-only arguments. PiperOrigin-RevId: 448066207 | 11 May 2022, 20:06:54 UTC |
| d092d63 | jax authors | 11 May 2022, 19:31:13 UTC | Merge pull request #10668 from sharadmv:custom-interpreter-update PiperOrigin-RevId: 448058058 | 11 May 2022, 19:31:13 UTC |
| 43467f9 | jax authors | 11 May 2022, 19:03:05 UTC | Merge pull request #10657 from sharadmv:effects-design-note PiperOrigin-RevId: 448051701 | 11 May 2022, 19:03:05 UTC |
| 590b916 | Peter Hawkins | 11 May 2022, 18:45:28 UTC | Add a sort_eigenvalues option to lax.linalg.eigh(). An upcoming change to add a more scalable QDWH-based TPU symmetric eigendecomposition requires that we can obtain the TPU eigenvalues unsorted. The option already exists in XLA, so we simply need to plumb it through to the lax primitive. PiperOrigin-RevId: 448047584 | 11 May 2022, 18:46:03 UTC |
| 2bce098 | Sharad Vikram | 10 May 2022, 17:43:03 UTC | Add sequencing effects design note | 11 May 2022, 18:19:58 UTC |
| 2867297 | Matthew Johnson | 11 May 2022, 18:18:25 UTC | fix grad(..., argnums=-1), regressed in #10453 | 11 May 2022, 18:19:22 UTC |
| aca9dc6 | Sharad Vikram | 11 May 2022, 18:07:00 UTC | Update custom interpreter tutorial | 11 May 2022, 18:11:47 UTC |
| 52ad3e6 | Tianjian Lu | 11 May 2022, 16:59:17 UTC | [linalg] Extend svd test coverage to input of zero and near-zero elements. PiperOrigin-RevId: 448019846 | 11 May 2022, 16:59:51 UTC |
| 1b0be50 | jax authors | 11 May 2022, 16:18:47 UTC | Merge pull request #10658 from ROCmSoftwarePlatform:rocm_unit_test_enablement PiperOrigin-RevId: 448010649 | 11 May 2022, 16:18:47 UTC |
| a076d34 | Filippo Vicentini | 11 May 2022, 15:57:26 UTC | Make CUDA install instructions copy-pasteablee | 11 May 2022, 15:57:26 UTC |
| 46d034b | Yash Katariya | 11 May 2022, 15:34:46 UTC | Add the nightly dev version to `__version__` of jaxlib. PiperOrigin-RevId: 448001375 | 11 May 2022, 15:35:16 UTC |
| 90f926a | Eugene Burmako | 10 May 2022, 22:38:41 UTC | [MHLO] Migrate mhlo.all_reduce to HLO_CompatibleOperandsAndResultType This runs into the currently unsupported feature in Python bindings which prevents it from taking advantage of the type inference functionality provided by HLO_CompatibleOperandsAndResultType. PiperOrigin-RevId: 447844880 | 10 May 2022, 22:39:18 UTC |
| dfb2caf | Yash Katariya | 10 May 2022, 21:05:09 UTC | Add nightly `__version__` string if building jaxlib nightly PiperOrigin-RevId: 447822974 | 10 May 2022, 21:05:35 UTC |
| 7d27343 | Sharad Vikram | 10 May 2022, 20:13:55 UTC | Attach keepalive to pmap executable PiperOrigin-RevId: 447810195 | 10 May 2022, 20:14:24 UTC |
| 48f47c3 | Tianjian Lu | 10 May 2022, 20:02:18 UTC | [linalg] Fix a bug in computing derivatives of scipy.special.lpmn. PiperOrigin-RevId: 447807140 | 10 May 2022, 20:02:58 UTC |
| 8d9f17d | Rohit Santhanam | 10 May 2022, 19:38:48 UTC | Disabled one and enabled several unit tests for ROCm. | 10 May 2022, 19:47:26 UTC |
| a62ca21 | Roy Frostig | 10 May 2022, 18:24:50 UTC | use `upgrade` option in defining our latest upgrade flag PiperOrigin-RevId: 447783540 | 10 May 2022, 18:25:23 UTC |
| b826328 | jax authors | 10 May 2022, 18:05:19 UTC | Merge pull request #10649 from onnoeberhard:patch-1 PiperOrigin-RevId: 447777762 | 10 May 2022, 18:05:19 UTC |
| 882a2d5 | Anselm Levskaya | 10 May 2022, 16:43:36 UTC | Rollback of PR #10393 "Improve performance of array integer indexing" This PR has broken some user models so needs to be investigated further before merging. PiperOrigin-RevId: 447756000 | 10 May 2022, 16:44:10 UTC |
| 7276731 | Onno Eberhard | 10 May 2022, 11:04:49 UTC | Small fix in FAQ documentation | 10 May 2022, 11:04:49 UTC |
| 68dbe74 | jax authors | 09 May 2022, 23:22:58 UTC | Merge pull request #10640 from jakevdp:update-myst-nb PiperOrigin-RevId: 447589472 | 09 May 2022, 23:22:58 UTC |
| 1fdc9a2 | Eugene Burmako | 09 May 2022, 22:44:33 UTC | [MHLO] Add CHLO lowering for top_k This follows the introduction of chlo.top_k in the diffbase. PiperOrigin-RevId: 447581682 | 09 May 2022, 22:45:04 UTC |
| 29efc30 | Jake VanderPlas | 09 May 2022, 21:32:48 UTC | DOC: update myst-nb to v0.15.0 | 09 May 2022, 21:32:48 UTC |
| bb56f40 | Matthew Johnson | 09 May 2022, 20:25:52 UTC | Internal change PiperOrigin-RevId: 447549479 | 09 May 2022, 20:26:30 UTC |
| 4bc1c1c | Tianjian Lu | 09 May 2022, 18:28:53 UTC | [linalg] Add svd on zero matrix. PiperOrigin-RevId: 447521398 | 09 May 2022, 18:29:22 UTC |
| 95b8e49 | jax authors | 09 May 2022, 18:07:25 UTC | Merge pull request #10636 from jakevdp:jupytext-doc PiperOrigin-RevId: 447515735 | 09 May 2022, 18:07:25 UTC |
| b5f10bc | Yash Katariya | 09 May 2022, 17:47:48 UTC | Fix the bug where allow_uneven_axes was set to True by default instead of reading from the function argument. PiperOrigin-RevId: 447510102 | 09 May 2022, 17:51:09 UTC |
| 9617391 | jax authors | 09 May 2022, 17:45:40 UTC | Add retry logic to `cloud_tpu_init.py` when API response fails PiperOrigin-RevId: 447509510 | 09 May 2022, 17:46:10 UTC |
| d3f89c1 | jax authors | 09 May 2022, 17:38:25 UTC | Merge pull request #10626 from ngam:patch-1 PiperOrigin-RevId: 447507358 | 09 May 2022, 17:38:25 UTC |
| fd66781 | jax authors | 09 May 2022, 17:03:48 UTC | Merge pull request #10573 from mattjj:simplify-var PiperOrigin-RevId: 447498112 | 09 May 2022, 17:03:48 UTC |
| 9679ca8 | Jake VanderPlas | 09 May 2022, 16:53:33 UTC | DOC: clarify jupytext instructions | 09 May 2022, 16:53:33 UTC |
| 705c07a | Matthew Johnson | 04 May 2022, 21:03:05 UTC | remove count attribute and total_ordering from core.Var Originally we used the 'Var.count' attribute to ensure Var instances were printed consistently regardless of context, even though only their object id was load-bearing. That is, Var.count was only used for pretty printing. (#1949 added a total_ordering on Var for reasons out of scope of JAX's core code.) But #8019 revised our pretty-printing so as not to use Var.count. Instead it chose how to pretty-print Var instances based on their order of appearance in a jaxpr. That meant Var.count really wasn't useful anymore. So this PR removes Var.count. In fact, Var.__repr__ and JaxprEqn.__repr__ were made confusing after #8019, since they could print variable names totally different from the names that would appear when the same JaxprEqn or Var objects were printed as part of a jaxpr. That is, before this PR< we might have a jaxpr which printed like: ```python import jax def f(x): for _ in range(3): x = jax.numpy.sin(x) return x jaxpr = jax.make_jaxpr(f)(3.) print(jaxpr) _, eqn, _ = jaxpr.jaxpr.eqns print(eqn) ``` Notice the variable names in the equation pretty-print don't correspond to any in the jaxpr pretty-print! So this PR changes JaxprEqn.__repr__ and Var.__repr__ to show Var object ids. | 09 May 2022, 16:31:23 UTC |
| 7132fba | ngam | 08 May 2022, 12:50:29 UTC | minor edits to clarify pjit docs | 09 May 2022, 16:14:45 UTC |
| fa96d9b | jax authors | 09 May 2022, 15:57:29 UTC | Merge pull request #10631 from gnecula:tf_limit2 PiperOrigin-RevId: 447481414 | 09 May 2022, 15:57:29 UTC |
| a50f35e | jax authors | 09 May 2022, 15:48:58 UTC | Merge pull request #10620 from JeppeKlitgaard:python310 PiperOrigin-RevId: 447479542 | 09 May 2022, 15:48:58 UTC |
| 91583f9 | George Necula | 09 May 2022, 15:40:41 UTC | [jax2tf] Update jax2tf limitations. It seems that TF has made progress and more JAX primitives can now be converted correctly to TF. This PR updates the limitations. | 09 May 2022, 15:40:41 UTC |
| e94221b | Peter Hawkins | 09 May 2022, 15:14:56 UTC | [MHLO] Migrate GPU select_and_gather_add translation rule to MHLO. PiperOrigin-RevId: 447472634 | 09 May 2022, 15:15:33 UTC |
| 4fa4872 | jax authors | 09 May 2022, 13:57:31 UTC | Merge pull request #10629 from ROCmSoftwarePlatform:rocm_unit_test_fixes PiperOrigin-RevId: 447458217 | 09 May 2022, 13:57:31 UTC |
| 43fb062 | jax authors | 09 May 2022, 13:49:18 UTC | Merge pull request #10628 from gnecula:tf_limitations PiperOrigin-RevId: 447456302 | 09 May 2022, 13:49:18 UTC |
| a5c3853 | Rohit Santhanam | 09 May 2022, 13:36:48 UTC | Fix unit test failures on ROCm. | 09 May 2022, 13:36:48 UTC |
| 137fdb4 | George Necula | 09 May 2022, 12:45:36 UTC | [jax2tf] Update JAX limitations. JAX has made progress in coverage of primitives on TPU. This PR updates those limitations. | 09 May 2022, 13:28:02 UTC |
| ccccd8a | jax authors | 09 May 2022, 04:21:49 UTC | Merge pull request #9868 from YouJiacheng:patch-1 PiperOrigin-RevId: 447377010 | 09 May 2022, 04:21:49 UTC |
| a11f15e | Jeppe Klitgaard | 07 May 2022, 12:38:55 UTC | feat: officially support Python 3.10 | 07 May 2022, 12:43:12 UTC |
| 78fafc5 | jax authors | 07 May 2022, 01:28:05 UTC | Merge pull request #10612 from ROCmSoftwarePlatform:rocm_build_fixes PiperOrigin-RevId: 447111173 | 07 May 2022, 01:28:05 UTC |
| 1093559 | Tianjian Lu | 06 May 2022, 23:33:13 UTC | [linalg] Add matmul precision scope for svd. PiperOrigin-RevId: 447095391 | 06 May 2022, 23:33:50 UTC |
| bbdcec8 | Rohit Santhanam | 06 May 2022, 22:57:51 UTC | Fixes to enable JAX to build on ROCm. | 06 May 2022, 22:57:51 UTC |
| 7251956 | Peter Hawkins | 06 May 2022, 22:33:41 UTC | Remove unused jaxlib build targets. //jaxlib:lapack is unused, and once it is gone we can merge //jaxlib:mhlo_helpers into //jaxlib. PiperOrigin-RevId: 447084931 | 06 May 2022, 22:34:14 UTC |
| 244bc8c | jax authors | 06 May 2022, 22:26:53 UTC | Merge pull request #10393 from lgeiger:speedup-indexing PiperOrigin-RevId: 447083562 | 06 May 2022, 22:26:53 UTC |
| 883cf2b | Peter Hawkins | 06 May 2022, 21:50:54 UTC | Refactor custom call building code in jaxlib to use a helper function. Refactoring only, no functional changes intended. This should fix a jaxlib build issue on Windows: we only have one constructor of layouts, and it explicitly requests an int64 type. Fixes https://github.com/google/jax/issues/10474 PiperOrigin-RevId: 447076192 | 06 May 2022, 21:51:24 UTC |
| cec717d | Tom Hennigan | 06 May 2022, 20:53:49 UTC | Return wrapper type in pytree rule for TF's DictWrapper. Fixes #10586 PiperOrigin-RevId: 447063792 | 06 May 2022, 20:54:21 UTC |
| 08c3c2e | Peter Hawkins | 06 May 2022, 20:47:23 UTC | Split CUDA and HIP C++ code in jaxlib into separate directories. PiperOrigin-RevId: 447062506 | 06 May 2022, 20:48:00 UTC |
| 79f0f6c | Peter Hawkins | 06 May 2022, 19:59:40 UTC | Revert: Make concatenate allow concatenation on dynamic dimensions Using -1 as a dynamic dimension is an MHLO convention not a JAX-level convention. PiperOrigin-RevId: 447052267 | 06 May 2022, 20:00:16 UTC |
| dcfb546 | jax authors | 06 May 2022, 19:48:20 UTC | Merge pull request #10607 from jakevdp:jit-faq PiperOrigin-RevId: 447049741 | 06 May 2022, 19:48:20 UTC |
| c4836aa | Jake VanderPlas | 06 May 2022, 19:39:12 UTC | DOC: add FAQ entry on jit-compiling methods | 06 May 2022, 19:39:12 UTC |
| 452de3f | jax authors | 06 May 2022, 18:01:39 UTC | Merge pull request #10604 from google:yashk2810-patch-11 PiperOrigin-RevId: 447025482 | 06 May 2022, 18:01:39 UTC |
| 40c61cf | Yash Katariya | 06 May 2022, 17:54:38 UTC | Update workspace for jaxlib release for testing | 06 May 2022, 17:54:38 UTC |
| 562e27d | Peter Hawkins | 06 May 2022, 16:34:25 UTC | Merge remaining CUDA and ROCM Python code. Completes work started in https://github.com/google/jax/pull/10556 PiperOrigin-RevId: 447005344 | 06 May 2022, 16:35:01 UTC |
