| 9d1cbc7 | Skye Wanderman-Milne | 22 August 2023, 20:10:12 UTC | Default to PJRT TPU runtime instead of StreamExecutor on older jaxlibs. I messed up the forwards compat in https://github.com/google/jax/commit/3e50fea29edb6b78426dde511414429f2d2fddf8. The next jaxlib release won't need the env var at all, but jaxlib 0.4.14 and older still do. | 23 August 2023, 00:06:16 UTC |
| 27d19ee | jax authors | 22 August 2023, 17:48:16 UTC | Instrument a metric to measure the number of tasks using compilation cache in JAX -> PJRT. Create the metric '/jax/compilation_cache/tasks_using_original_cache' to record the number of tasks using compilation cache. PiperOrigin-RevId: 559159282 | 22 August 2023, 17:48:58 UTC |
| c1f6b9b | Chris Jones | 22 August 2023, 15:47:17 UTC | [jax_triton] Always use "triton_kernel_call" as custom call target name. Now there is a `name` field in the serialized proto, there is an alternative way to differentiate kernel calls. PiperOrigin-RevId: 559124166 | 22 August 2023, 15:48:02 UTC |
| a77da34 | jax authors | 22 August 2023, 13:14:16 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/e3533af05fac51df26322c89e12750da99936084. PiperOrigin-RevId: 559090025 | 22 August 2023, 13:15:25 UTC |
| 8891503 | George Necula | 22 August 2023, 10:46:38 UTC | [callback] Add workaround for TPU host_callback not supporting empty arrays. Currently JAX callbacks on TPU raise errors when the called function takes empty arguments or returns empty results. It seems that the send_to_host function works even with empty arrays, but recv_from_host crashes (crash log below). Here we work around this issue, by ensuring that only the non-empty results of the Python callback are sent to the device computation and the empty results are replaced with empty constants in the device computation. This is part of the work to replace uses of host_callback with io_callback. PiperOrigin-RevId: 559061336 | 22 August 2023, 10:47:18 UTC |
| f70f1f8 | Chris Jones | 22 August 2023, 10:04:40 UTC | Internal change. PiperOrigin-RevId: 559053761 | 22 August 2023, 10:05:17 UTC |
| 00017c7 | jax authors | 22 August 2023, 01:49:01 UTC | Merge pull request #17204 from jakevdp:norm-sf PiperOrigin-RevId: 558959421 | 22 August 2023, 01:49:01 UTC |
| f2d7798 | jax authors | 21 August 2023, 23:53:48 UTC | Add `register_static` decorator to tree_util to facilitate creating leafless classes. PiperOrigin-RevId: 558937697 | 21 August 2023, 23:54:29 UTC |
| cf11f8d | Jake VanderPlas | 21 August 2023, 23:48:39 UTC | stats.norm: add logsf & make sf more accurate near zero | 21 August 2023, 23:48:39 UTC |
| 9f5999d | Peter Hawkins | 21 August 2023, 22:55:02 UTC | Improve type annotations for jax.numpy. * Allow sequences of axes to jnp.flip, rather than mandating tuples. Users sometimes pass lists here. * Allow array-like pad_width values to pad(). PiperOrigin-RevId: 558923802 | 21 August 2023, 22:56:14 UTC |
| d03187e | jax authors | 21 August 2023, 20:37:24 UTC | Merge pull request #17060 from nouiz:rosetta_doc PiperOrigin-RevId: 558886922 | 21 August 2023, 20:37:24 UTC |
| 2f28848 | jax authors | 21 August 2023, 19:47:45 UTC | Merge pull request #17202 from jakevdp:logsumexp-types PiperOrigin-RevId: 558873117 | 21 August 2023, 19:47:45 UTC |
| 27324ff | Jake VanderPlas | 21 August 2023, 16:10:19 UTC | special.logsumexp: fix incorrect annotation | 21 August 2023, 16:10:19 UTC |
| a14d64b | jax authors | 21 August 2023, 13:03:54 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/c1e30de7d34a13e6e2c3dcfd4906b4255c98c99e. PiperOrigin-RevId: 558765656 | 21 August 2023, 13:04:36 UTC |
| aed80a7 | Yash Katariya | 21 August 2023, 01:56:50 UTC | Add docs for `is_fully_addressable` to jax.Array and remove GDA from process_allgather docs and clarify it a bit more. PiperOrigin-RevId: 558643985 | 21 August 2023, 01:58:22 UTC |
| b4a6284 | jax authors | 19 August 2023, 12:25:34 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/51c0ad464e05c18098df7b771a7e02e1ad86dd27. PiperOrigin-RevId: 558390333 | 19 August 2023, 12:26:07 UTC |
| 2eb9a6a | Sharad Vikram | 19 August 2023, 05:57:02 UTC | Disable test_set_vmap on GPU due to flakiness PiperOrigin-RevId: 558336356 | 19 August 2023, 05:57:54 UTC |
| a8945fd | jax authors | 19 August 2023, 02:59:26 UTC | Add an unregister event listener function in JAX monitoring. Add a private function _unregister_event_listener_by_callback to remove registered event listeners. The functions are supposed to be called in test only. Add a getter function for event listeners to help unit testing the unregister function. PiperOrigin-RevId: 558309557 | 19 August 2023, 03:00:03 UTC |
| 1a9109f | jax authors | 18 August 2023, 22:06:04 UTC | Merge pull request #17184 from hawkinsp:devicearray PiperOrigin-RevId: 558257332 | 18 August 2023, 22:06:04 UTC |
| 2c32660 | Peter Hawkins | 18 August 2023, 20:50:36 UTC | Replace references to DeviceArray with Array. A number of stale references are lurking in our documentation. | 18 August 2023, 21:46:00 UTC |
| 97af33c | jax authors | 18 August 2023, 21:29:22 UTC | Merge pull request #17182 from hawkinsp:tril PiperOrigin-RevId: 558247039 | 18 August 2023, 21:29:22 UTC |
| a80cbc5 | Skye Wanderman-Milne | 18 August 2023, 21:19:49 UTC | [JAX] Implement the `stream` argument to jax.Array.__dlpack__ for CUDA GPU Also implements jax.Array.__dlpack_device__. See https://dmlc.github.io/dlpack/latest/python_spec.html This requires plumbing the raw CUDA stream pointer through PJRT and StreamExecutor (since the GPU PJRT implementation is still based on SE). This is done via the new PJRT method ExternalReference::WaitUntilBufferReadyOnStream. I haven't plumbed this through the PJRT C API yet, because I'm still debating whether this should be part of the main API or a GPU-specific extension (plus either way it should probably be its own change). PiperOrigin-RevId: 558245360 | 18 August 2023, 21:20:38 UTC |
| 3119b43 | jax authors | 18 August 2023, 20:45:42 UTC | Merge pull request #17181 from hawkinsp:docs PiperOrigin-RevId: 558236379 | 18 August 2023, 20:45:42 UTC |
| 4224a4d | Peter Hawkins | 18 August 2023, 20:14:42 UTC | Deprecate jax.scipy.linalg.tril and jax.scipy.linalg.triu. The corresponding functions are deprecated in scipy. Use the equivalent jax.numpy functions instead. | 18 August 2023, 20:14:42 UTC |
| 5d680d4 | Peter Hawkins | 18 August 2023, 19:56:35 UTC | Improve jax.nn documentation. Fixes https://github.com/google/jax/issues/17171 | 18 August 2023, 19:56:35 UTC |
| 9d908d9 | jax authors | 18 August 2023, 18:42:05 UTC | Merge pull request #17169 from eltociear:eltociear-patch-1 PiperOrigin-RevId: 558205014 | 18 August 2023, 18:42:05 UTC |
| 4af1cd6 | jax authors | 18 August 2023, 18:22:40 UTC | Merge pull request #17164 from 8bitmp3:jax-readme-update PiperOrigin-RevId: 558199294 | 18 August 2023, 18:22:40 UTC |
| 6d68613 | jax authors | 18 August 2023, 18:12:18 UTC | Force inline all the calls in the module created by a custom call. PiperOrigin-RevId: 558196979 | 18 August 2023, 18:13:45 UTC |
| 209b6b0 | jax authors | 18 August 2023, 18:04:43 UTC | Merge pull request #17144 from jakevdp:zeta PiperOrigin-RevId: 558193896 | 18 August 2023, 18:04:43 UTC |
| 2e13243 | Yash Katariya | 18 August 2023, 17:54:03 UTC | Check that jaxlib's nightly date is not greater than jax's nightly date in `check_jaxlib_version` PiperOrigin-RevId: 558191583 | 18 August 2023, 17:54:52 UTC |
| e01f261 | Yash Katariya | 18 August 2023, 15:32:32 UTC | Add jax release jobs. Also make the `__version__` be the `.dev` version by default unless `JAX_RELEASE` env var is set. PiperOrigin-RevId: 558153979 | 18 August 2023, 15:33:06 UTC |
| 3323692 | jax authors | 18 August 2023, 12:34:24 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/465c945d598fc006b8e96725dd34eb0e7938a863. PiperOrigin-RevId: 558119555 | 18 August 2023, 12:34:56 UTC |
| edb88ce | George Necula | 18 August 2023, 07:57:01 UTC | [jax2tf] Disable a few checks on GPU also, for mixed-precision dot_general. These tests were already disabled on CPU and TPU, but it seems these are due to general limitations of TF/XLA, including on GPU. PiperOrigin-RevId: 558066104 | 18 August 2023, 07:57:40 UTC |
| eaa2712 | Ikko Eltociear Ashimine | 18 August 2023, 07:00:38 UTC | Fix typo in hlo_helpers.py libary -> library | 18 August 2023, 07:00:38 UTC |
| 6868be3 | Yash Katariya | 18 August 2023, 02:02:24 UTC | Create a jax nightly wheel and upload to GCS. PiperOrigin-RevId: 557998568 | 18 August 2023, 02:02:59 UTC |
| 6d1b5d9 | jax authors | 18 August 2023, 00:42:11 UTC | Merge pull request #17160 from jakevdp:dep-issubsctype PiperOrigin-RevId: 557984634 | 18 August 2023, 00:42:11 UTC |
| d300c5c | 8bitmp3 | 17 August 2023, 23:42:19 UTC | Update JAX README.md - remove jax.example_libraries, stax | 17 August 2023, 23:42:19 UTC |
| f076197 | jax authors | 17 August 2023, 22:06:11 UTC | Merge pull request #17163 from jakevdp:binom-doc PiperOrigin-RevId: 557947436 | 17 August 2023, 22:06:11 UTC |
| cc773d4 | Jake VanderPlas | 17 August 2023, 19:56:15 UTC | Add missing jax.scipy.stats docs | 17 August 2023, 19:56:15 UTC |
| 8bba992 | Jake VanderPlas | 17 August 2023, 19:27:52 UTC | deprecate jax.numpy.issubsctype | 17 August 2023, 19:27:52 UTC |
| 2f6dec1 | jax authors | 17 August 2023, 16:45:22 UTC | Merge pull request #17152 from gnecula:tf_dot_general PiperOrigin-RevId: 557851430 | 17 August 2023, 16:45:22 UTC |
| d0059b7 | Frederic Bastien | 17 August 2023, 14:46:35 UTC | typo | 17 August 2023, 14:46:35 UTC |
| 8894892 | Peter Hawkins | 17 August 2023, 13:43:31 UTC | Remove the `canonicalize_dtypes` argument from mlir.ir_constant(s). Instead, force the caller to explicitly canonicalize the argument if that's what they want. The current behavior (canonicalize by default) is not the behavior we want to encourage: we want to canonicalize exactly where we need to and nowhere else. PiperOrigin-RevId: 557806903 | 17 August 2023, 13:44:12 UTC |
| ab9555e | jax authors | 17 August 2023, 13:11:05 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/b55511f5fa76a5cf83bd8912ed34813f8b40888b. PiperOrigin-RevId: 557800589 | 17 August 2023, 13:19:19 UTC |
| e2ebd9a | jax authors | 17 August 2023, 12:01:50 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/c7e022af036940588d8fbc097155c21aa53cdbf7. PiperOrigin-RevId: 557787256 | 17 August 2023, 12:02:31 UTC |
| a833652 | jax authors | 17 August 2023, 11:23:02 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/c782fe83564beefc7341dea151fe57b4ec803094. PiperOrigin-RevId: 557780339 | 17 August 2023, 11:35:07 UTC |
| 98c5eab | jax authors | 17 August 2023, 11:17:28 UTC | Merge pull request #17154 from gnecula:poly_max_error_inputs PiperOrigin-RevId: 557778341 | 17 August 2023, 11:17:28 UTC |
| 1096dca | jax authors | 17 August 2023, 11:10:49 UTC | Fix a typo. PiperOrigin-RevId: 557778263 | 17 August 2023, 11:17:06 UTC |
| 5ab3e9f | jax authors | 17 August 2023, 11:06:28 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/2a2bc3d9438132cc1f263d20f49a7fb48a8db9be. PiperOrigin-RevId: 557777322 | 17 August 2023, 11:07:20 UTC |
| 26eb6c3 | George Necula | 17 August 2023, 09:19:04 UTC | [shape_poly] Relax the limit on the number of error inputs for shape assertions Lift the limit from 4 to 32 to follow the change in tf.XlaCallModule of this limit. Currently, the error message formatter needs at most 6 error message inputs. | 17 August 2023, 10:15:05 UTC |
| 4c0664a | jax authors | 17 August 2023, 10:11:22 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/f8dfd1ee600ff21b34db27b8bb7417ce2ab4a884. PiperOrigin-RevId: 557766683 | 17 August 2023, 10:11:51 UTC |
| 4f3d9df | jax authors | 17 August 2023, 09:54:39 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/08af1e149d1606b495866f3d8f51bba69d494be3. PiperOrigin-RevId: 557763125 | 17 August 2023, 09:55:13 UTC |
| 49667e5 | jax authors | 17 August 2023, 09:37:58 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/4acd835d9c48ff86559791ae3d683de23607053e. PiperOrigin-RevId: 557759271 | 17 August 2023, 09:38:38 UTC |
| 18e5e9b | George Necula | 17 August 2023, 09:05:16 UTC | [jax2tf] Fix for dot_general with different dtypes for graph serialization. Graph serialization did not work for dot_general with different dtypes. Added a workaround by converting the operands to the smallest type that can represent both operands. This is not accurate w.r.t. what JAX native does (use native serialization for that), but avoids failures. Added tests for these cases for enable_xla=[True|False]. Some of these tests fail, e.g., when involving x64 types for enable_xla=True and when involving certain unsupported types for enable_xla=False (uint16, int16). | 17 August 2023, 09:26:59 UTC |
| 7a5b9ff | jax authors | 17 August 2023, 04:57:43 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/225f05fc6d792b4e39a42e588f679e0b747f515e. PiperOrigin-RevId: 557700420 | 17 August 2023, 04:58:34 UTC |
| fcc44fb | jax authors | 17 August 2023, 04:41:08 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/3bbf0a04496c636e85a2bf5e02edbfef1f080fa6. PiperOrigin-RevId: 557696920 | 17 August 2023, 04:41:49 UTC |
| 1be198c | jax authors | 17 August 2023, 04:21:14 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/02d59d5d9a59207da44f7bc4c0f83309aac25046. PiperOrigin-RevId: 557693150 | 17 August 2023, 04:21:46 UTC |
| 0045b03 | jax authors | 17 August 2023, 03:59:09 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/1e81d47ecb624cb8e75890fe9c1e5b04558ef545. PiperOrigin-RevId: 557689198 | 17 August 2023, 03:59:47 UTC |
| 86380b0 | jax authors | 17 August 2023, 03:41:20 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/c3512249ed63ae65769be6874de98497b848e3d2. PiperOrigin-RevId: 557686293 | 17 August 2023, 03:42:01 UTC |
| 82e840c | jax authors | 17 August 2023, 02:17:23 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/ac00a88b0bbed0f6b1536d6cd52083d4d74ef710. PiperOrigin-RevId: 557671298 | 17 August 2023, 02:18:00 UTC |
| ebd09cc | jax authors | 17 August 2023, 02:02:33 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/d79398ede04347dde1b43b76249155ce481c7c7f. PiperOrigin-RevId: 557668919 | 17 August 2023, 02:03:24 UTC |
| 9c4ba21 | jax authors | 17 August 2023, 01:46:29 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/4ca5589bff25a944a556c9dba1b228638e8f7c19. PiperOrigin-RevId: 557666310 | 17 August 2023, 01:50:41 UTC |
| 242c2c1 | Yash Katariya | 17 August 2023, 01:40:53 UTC | Use `_internal_device_list` in __hash__ and __eq__ of Shardings and Mesh to speed them up. PiperOrigin-RevId: 557665385 | 17 August 2023, 01:41:58 UTC |
| 03a0d73 | jax authors | 17 August 2023, 00:27:59 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/28e756940622ba628dcd11b3a4dcbb1206d8f1f7. PiperOrigin-RevId: 557651441 | 17 August 2023, 00:28:43 UTC |
| 124bdb7 | jax authors | 17 August 2023, 00:12:37 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/51f7513132a2025629d294488643a7e85f90e22b. PiperOrigin-RevId: 557647957 | 17 August 2023, 00:13:41 UTC |
| 785af82 | Sharad Vikram | 17 August 2023, 00:04:28 UTC | [Pallas] Add Mosaic lowering for exp2 PiperOrigin-RevId: 557646023 | 17 August 2023, 00:05:03 UTC |
| caee312 | Sharad Vikram | 16 August 2023, 23:48:36 UTC | Expose exp2_p in jax.lax PiperOrigin-RevId: 557642106 | 16 August 2023, 23:49:15 UTC |
| 69cf0da | jax authors | 16 August 2023, 23:37:07 UTC | Merge pull request #17149 from jakevdp:typing-fix PiperOrigin-RevId: 557638893 | 16 August 2023, 23:37:07 UTC |
| 3bffa07 | jax authors | 16 August 2023, 23:14:40 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/13bb9d081295fe3d80350dd47d65e13e339b3b91. PiperOrigin-RevId: 557633286 | 16 August 2023, 23:15:51 UTC |
| c6f40e2 | Parker Schuh | 16 August 2023, 22:57:04 UTC | Reverts 75c3457264f9cc117ff09551ce3174d72689fa3d PiperOrigin-RevId: 557628297 | 16 August 2023, 23:06:28 UTC |
| c91f010 | jax authors | 16 August 2023, 22:56:33 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/4c2bf94fbbe6a78b46ccdb207b40a5495108f7a6. PiperOrigin-RevId: 557628169 | 16 August 2023, 22:57:11 UTC |
| 329791a | jax authors | 16 August 2023, 22:40:42 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/c05021a5bb58fbdb9966495bcfefd37fd098aae5. PiperOrigin-RevId: 557623875 | 16 August 2023, 22:41:24 UTC |
| e959c04 | jax authors | 16 August 2023, 22:22:25 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/6785669435c98591e789f1509d3c60cf94175741. PiperOrigin-RevId: 557618963 | 16 August 2023, 22:23:11 UTC |
| 9aca944 | Jake VanderPlas | 16 August 2023, 22:07:48 UTC | Fix type annotation for tree_util.default_registry | 16 August 2023, 22:07:48 UTC |
| bb4daa3 | jax authors | 16 August 2023, 21:58:25 UTC | Merge pull request #17147 from jakevdp:mypy-15 PiperOrigin-RevId: 557611626 | 16 August 2023, 21:58:25 UTC |
| 42d5441 | jax authors | 16 August 2023, 21:45:31 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/b3f12d680d9876ca6197d0d8d1aedadce9f50ca4. PiperOrigin-RevId: 557607869 | 16 August 2023, 21:46:11 UTC |
| 1923e2c | jax authors | 16 August 2023, 21:31:52 UTC | Merge pull request #17142 from jakevdp:polygamma PiperOrigin-RevId: 557603828 | 16 August 2023, 21:31:52 UTC |
| 6097d8f | Jake VanderPlas | 16 August 2023, 21:26:04 UTC | CI: bump mypy to version 1.5.0 | 16 August 2023, 21:26:04 UTC |
| 6cd467f | Jake VanderPlas | 16 August 2023, 20:43:41 UTC | Create lax.zeta with native HLO lowering | 16 August 2023, 20:43:41 UTC |
| 8d444a5 | jax authors | 16 August 2023, 20:36:15 UTC | Merge pull request #17145 from jakevdp:mypy-fix PiperOrigin-RevId: 557587526 | 16 August 2023, 20:36:15 UTC |
| f8a4afd | Jake VanderPlas | 16 August 2023, 20:27:11 UTC | Fix mypy error on Python 3.9 | 16 August 2023, 20:27:11 UTC |
| 4cf85b9 | jax authors | 16 August 2023, 19:12:16 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/baaeccccf8a205ede61be49652237b857907147a. PiperOrigin-RevId: 557563190 | 16 August 2023, 19:13:02 UTC |
| 0ad6196 | Jake VanderPlas | 16 August 2023, 18:57:05 UTC | Create lax.polygamma with native HLO lowering | 16 August 2023, 18:57:05 UTC |
| 539a6d1 | jax authors | 16 August 2023, 18:34:17 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/0cebcdd469f829f182424b85f8bbe9a188e19344. PiperOrigin-RevId: 557551693 | 16 August 2023, 18:35:08 UTC |
| a3cd351 | jax authors | 16 August 2023, 17:45:40 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/af311aff5c7d02801d2023e7809d046aa43bcc51. PiperOrigin-RevId: 557535177 | 16 August 2023, 17:46:45 UTC |
| b4bd1da | jax authors | 16 August 2023, 17:29:54 UTC | Merge pull request #17141 from 8bitmp3:jax-faq-change PiperOrigin-RevId: 557528004 | 16 August 2023, 17:29:54 UTC |
| 6eef465 | jax authors | 16 August 2023, 17:19:29 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/cd236825023fd7cda737f1cd5a8e0812cf1ec035. PiperOrigin-RevId: 557526553 | 16 August 2023, 17:29:39 UTC |
| a1d8787 | jax authors | 16 August 2023, 17:01:49 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/2784a5888082be4b8ae312aa901a9c3898a9af91. PiperOrigin-RevId: 557520885 | 16 August 2023, 17:10:45 UTC |
| ad15a38 | George Necula | 16 August 2023, 17:01:11 UTC | [host_callback] Remove old backwards compatibility flag jax_host_callback_ad_transforms. This flag was added in https://github.com/google/jax/pull/8678 in December 2021 when we changed the behavior of host_callback to not have special handling for autodiff. Nobody is using that flag now. This is part of a longer project to replace uses of host_callback with jax.pure_callback and jax.experimental.io_callback. PiperOrigin-RevId: 557520668 | 16 August 2023, 17:01:49 UTC |
| 14fa067 | jax authors | 16 August 2023, 16:45:08 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/26815644192344a3d076ad93b2ef5d83774fca6e. PiperOrigin-RevId: 557515766 | 16 August 2023, 16:46:21 UTC |
| cff08d1 | Frederic Bastien | 16 August 2023, 16:21:55 UTC | Small doc update following code review | 16 August 2023, 16:23:58 UTC |
| 05c7438 | Yash Katariya | 16 August 2023, 16:12:15 UTC | Skip PallasTest when triton is not installed. PiperOrigin-RevId: 557506520 | 16 August 2023, 16:12:51 UTC |
| 75c3457 | jax authors | 16 August 2023, 15:33:00 UTC | Merge pull request #16733 from jakevdp:einsum-mixed-precision PiperOrigin-RevId: 557492885 | 16 August 2023, 15:33:00 UTC |
| 6405901 | Jake VanderPlas | 16 August 2023, 15:22:58 UTC | Re-enable mixed-precision testDot on TPU The underlying issue has been fixed in XLA PiperOrigin-RevId: 557492724 | 16 August 2023, 15:23:58 UTC |
| f1a2bbf | Peter Hawkins | 16 August 2023, 15:08:51 UTC | Update XLA dependency to use revision http://github.com/openxla/xla/commit/2c817aeb83dcdab68dca36c0598d62cfbadda960. PiperOrigin-RevId: 557489343 | 16 August 2023, 15:09:24 UTC |
| 4ac2bdc | Chris Jones | 16 August 2023, 09:53:05 UTC | [jax_triton] Add user-specified `name` field to serialized format. PiperOrigin-RevId: 557415723 | 16 August 2023, 09:53:51 UTC |
| 14d52fc | Jake VanderPlas | 21 July 2023, 19:22:22 UTC | jnp.einsum: lower to mixed-precision dot_general when possible | 15 August 2023, 22:57:19 UTC |
| c7e8b81 | Jieying Luo | 15 August 2023, 22:23:17 UTC | [PJRT C API] Let framework explicitly check whether a plugin is initialized and initialize the plugin. Before this change, PJRT_Plugin_Initialize was called in LoadPjrtPlugin, which is only used in dynamic linking case. This change adds a bool and a method to check whether the plugin is initialized. The framework will explicitly check whether a plugin is initialized, and call InitializePjrtPlugin if it is not. This will be apply to both static linking and dynamic linking case. PiperOrigin-RevId: 557268670 | 15 August 2023, 22:24:14 UTC |
| b779671 | Yash Katariya | 15 August 2023, 20:25:37 UTC | Only donate if memory kinds match. This shouldn't break any existing behavior. Also only warn about unused donations if there are unused donations. PiperOrigin-RevId: 557231080 | 15 August 2023, 20:26:18 UTC |
| 6cabc36 | 8bitmp3 | 15 August 2023, 20:09:09 UTC | Update JAX F.A.Q. | 15 August 2023, 20:09:09 UTC |
