| e856ecf | jax authors | 23 February 2022, 05:12:39 UTC | Introduce a new PjRtEvent class, and move BlockHostUntilReady and OnReady methods from PjRtBuffer into PjRtEvent, so instead of calling buf->BlockHostUntilReady() you now call buf->GetEvent()->BlockHostUntilReady(). The intention is that all PjRt methods that currently enqueue future work but don't return a PjRtBuffer will be modified to return a PjRtEvent, and that all blocking and callbacks on futures will be performed via PjRtEvent. PiperOrigin-RevId: 430363082 | 03 March 2022, 00:03:18 UTC |
| 790b4bc | Yash Katariya | 02 March 2022, 22:25:33 UTC | Fix a typo (titled -> tiled) PiperOrigin-RevId: 432028664 | 02 March 2022, 22:26:08 UTC |
| 38ea085 | Jake VanderPlas | 02 March 2022, 20:18:09 UTC | lax_numpy.py: factor ndarray into its own module This is the first step in relanding the larger refactoring in #9724, which had to be rolled back due to downstream breakages. PiperOrigin-RevId: 431999528 | 02 March 2022, 20:18:46 UTC |
| 98572a6 | jax authors | 02 March 2022, 19:33:07 UTC | Merge pull request #9737 from jakevdp:cleanup-constant-like PiperOrigin-RevId: 431988229 | 02 March 2022, 19:33:07 UTC |
| f1e71c1 | jax authors | 02 March 2022, 18:11:09 UTC | [Jax] Format ann docstring. PiperOrigin-RevId: 431968329 | 02 March 2022, 18:11:52 UTC |
| 00e040e | Jake VanderPlas | 02 March 2022, 17:13:58 UTC | cleanup: remove _constant_like in favor of lax._const | 02 March 2022, 17:13:58 UTC |
| fb44d7c | jax authors | 02 March 2022, 16:58:22 UTC | [JAX] Add release note for the graduration of the experimental.ann module. PiperOrigin-RevId: 431951602 | 02 March 2022, 16:58:53 UTC |
| 3766dd2 | jax authors | 02 March 2022, 03:38:50 UTC | Rollback of: d09d7b8d1363eab1c14051eb2376e605366537f9 by Jake VanderPlas <jakevdp@google.com>: Factor-out pieces of lax_numpy.py PiperOrigin-RevId: 431833044 | 02 March 2022, 03:39:31 UTC |
| 4755dc3 | Yash Katariya | 02 March 2022, 01:44:14 UTC | Don't compute global shards if local_mesh == global_mesh and just return local_shards. PiperOrigin-RevId: 431815962 | 02 March 2022, 01:48:29 UTC |
| 72cc567 | Yash Katariya | 02 March 2022, 01:43:11 UTC | Use the new `mesh` property instead of the private `_global_mesh` attribute. PiperOrigin-RevId: 431815802 | 02 March 2022, 01:43:40 UTC |
| d9f82f7 | jax authors | 01 March 2022, 22:46:04 UTC | [JAX] Move `experimental.ann.approx_*_k` into `lax`. Updated docs, tests and the example code snippets. PiperOrigin-RevId: 431781401 | 01 March 2022, 22:46:33 UTC |
| 1a1bf12 | jax authors | 01 March 2022, 22:24:21 UTC | [JAX] Fix github lint_and_typecheck for Primitives PiperOrigin-RevId: 431776305 | 01 March 2022, 22:24:59 UTC |
| c50b9d4 | jax authors | 01 March 2022, 21:39:09 UTC | Merge pull request #9736 from mattjj:test-issue-9731 PiperOrigin-RevId: 431765284 | 01 March 2022, 21:39:09 UTC |
| 4075b81 | Matthew Johnson | 01 March 2022, 21:10:28 UTC | add regression test for #9731 | 01 March 2022, 21:10:28 UTC |
| d6812b2 | jax authors | 01 March 2022, 20:50:17 UTC | Merge pull request #9732 from jakevdp:jnp-copy PiperOrigin-RevId: 431754341 | 01 March 2022, 20:50:17 UTC |
| ed25509 | Jake VanderPlas | 01 March 2022, 18:44:33 UTC | implement jnp.copy | 01 March 2022, 19:56:36 UTC |
| 11664f8 | Yash Katariya | 01 March 2022, 19:45:26 UTC | Make the global_mesh that GDA has public PiperOrigin-RevId: 431739364 | 01 March 2022, 19:45:54 UTC |
| 5b30988 | jax authors | 01 March 2022, 19:33:39 UTC | Merge pull request #9724 from jakevdp:refactor-lax-numpy PiperOrigin-RevId: 431736367 | 01 March 2022, 19:33:39 UTC |
| d09d7b8 | Jake VanderPlas | 01 March 2022, 17:40:37 UTC | Factor-out pieces of lax_numpy.py | 01 March 2022, 17:40:37 UTC |
| cffe997 | Peter Hawkins | 01 March 2022, 16:55:29 UTC | Handle jaxpr constants correctly in MLIR lowering of conditional branches. Add some dynamic type checks and type annotations to catch this kind of problem sooner. There's no test case, because I'm not entirely sure how to make a test case for this. In fact, I'm not even sure it's legal for a conditional branch to have non-empty constants. We'll dig into that separately. PiperOrigin-RevId: 431697808 | 01 March 2022, 16:56:08 UTC |
| c7508d1 | jax authors | 28 February 2022, 20:59:08 UTC | Merge pull request #9721 from jakevdp:poisson-nan PiperOrigin-RevId: 431505317 | 28 February 2022, 20:59:08 UTC |
| d0cc339 | Yash Katariya | 28 February 2022, 20:53:57 UTC | Add `block_until_ready` method to GDA PiperOrigin-RevId: 431504594 | 28 February 2022, 20:54:23 UTC |
| 2c2773a | Jake VanderPlas | 28 February 2022, 20:10:47 UTC | jax.random.poisson: fix corner cases | 28 February 2022, 20:10:47 UTC |
| 92cb865 | jax authors | 28 February 2022, 19:54:29 UTC | Merge pull request #9718 from google:dependabot/github_actions/actions/setup-python-3 PiperOrigin-RevId: 431490872 | 28 February 2022, 19:54:29 UTC |
| c339330 | Peter Hawkins | 28 February 2022, 17:26:20 UTC | [XLA:CPU] Relax test tolerances for tests using XLA:CPU. An upcoming change to XLA:CPU will disable reassociation on floating point operators by default which is an unsound fast math optimization. This change is being made to fix numerical errors in softmax computations caused by reassocation. After that change, we will enable reassociation only in reduction operators where it is very important for performance and the XLA operator contract allows that. Since this change alters the order of operations, it may cause small numerical changes leading to test failures. This change relaxes test tolerances to make tests pass. PiperOrigin-RevId: 431453240 | 28 February 2022, 17:26:54 UTC |
| 680c06d | dependabot[bot] | 28 February 2022, 17:11:41 UTC | Bump actions/setup-python from 2 to 3 Bumps [actions/setup-python](https://github.com/actions/setup-python) from 2 to 3. - [Release notes](https://github.com/actions/setup-python/releases) - [Commits](https://github.com/actions/setup-python/compare/v2...v3) --- updated-dependencies: - dependency-name: actions/setup-python dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] <support@github.com> | 28 February 2022, 17:11:41 UTC |
| 043561a | jax authors | 28 February 2022, 16:57:42 UTC | Merge pull request #9711 from jblespiau:changelist/431385993 PiperOrigin-RevId: 431447036 | 28 February 2022, 16:57:42 UTC |
| 0b9ae98 | jax authors | 28 February 2022, 15:56:57 UTC | Merge pull request #9701 from jakevdp:fix-sharp-bits PiperOrigin-RevId: 431435194 | 28 February 2022, 15:56:57 UTC |
| 25472c2 | Jean-Baptiste Lespiau | 28 February 2022, 10:32:32 UTC | Remove an unnecessary condition. This makes sure jit, xla_computation and pjit share the same logic for processing static arguments. | 28 February 2022, 10:34:01 UTC |
| c66fbeb | jax authors | 26 February 2022, 07:29:04 UTC | Merge pull request #9193 from froystig:xla-executable-attr PiperOrigin-RevId: 431102224 | 26 February 2022, 07:29:04 UTC |
| d636e74 | Roy Frostig | 13 January 2022, 23:42:17 UTC | make `xla_executable` a property, consistent across executable types Also test IR and executable-related methods of `Lowered` and `Compiled`. | 26 February 2022, 03:05:44 UTC |
| d583f87 | jax authors | 25 February 2022, 22:19:17 UTC | Merge pull request #9706 from samuela:samuela/mkl-fix PiperOrigin-RevId: 431028480 | 25 February 2022, 22:19:17 UTC |
| bf59b7d | Samuel Ainsworth | 25 February 2022, 22:02:55 UTC | Relax tolerances slightly for MKL. Fix https://github.com/google/jax/issues/9705. | 25 February 2022, 22:02:55 UTC |
| 7ec04c8 | jax authors | 25 February 2022, 21:23:11 UTC | Merge pull request #9702 from jakevdp:fix-load PiperOrigin-RevId: 431016627 | 25 February 2022, 21:23:11 UTC |
| 1b01865 | Jake VanderPlas | 25 February 2022, 17:27:42 UTC | BUG: return numpy arrays for jnp.load() with unsupported dtypes | 25 February 2022, 17:27:42 UTC |
| 74e7fdf | Jake VanderPlas | 25 February 2022, 17:15:55 UTC | DOC: update sharp bits info about fori_loop differentiability | 25 February 2022, 17:15:55 UTC |
| d5a1c64 | jax authors | 25 February 2022, 01:51:09 UTC | Merge pull request #9623 from rsuderman:FixIreeJax PiperOrigin-RevId: 430834942 | 25 February 2022, 01:51:09 UTC |
| 6914bae | jax authors | 25 February 2022, 01:13:48 UTC | Merge pull request #9641 from rsuderman:FixDefaultDevice PiperOrigin-RevId: 430828922 | 25 February 2022, 01:13:48 UTC |
| ffcf477 | jax authors | 25 February 2022, 01:07:02 UTC | Merge pull request #9684 from jakevdp:fix-changelog PiperOrigin-RevId: 430827769 | 25 February 2022, 01:07:02 UTC |
| 5172703 | Jake VanderPlas | 24 February 2022, 16:18:30 UTC | Remove duplicate changelog entry | 24 February 2022, 16:18:30 UTC |
| d2834b6 | jax authors | 24 February 2022, 16:12:41 UTC | Merge pull request #9658 from tornikeo:docs-copybutton PiperOrigin-RevId: 430705653 | 24 February 2022, 16:12:41 UTC |
| 8372b98 | jax authors | 24 February 2022, 15:22:52 UTC | [JAX] Move ann.ann_recall back to tests. The function is simple enough for users to implement their own on the host. PiperOrigin-RevId: 430696789 | 24 February 2022, 15:23:17 UTC |
| cdd0c68 | jax authors | 24 February 2022, 15:02:34 UTC | Merge pull request #9619 from hawkinsp:index2 PiperOrigin-RevId: 430692755 | 24 February 2022, 15:02:34 UTC |
| f2c1743 | jax authors | 24 February 2022, 14:57:24 UTC | [JAX] Update comments and documents for ANN. Copied some of the examples from the module to the functions to make the migration to jax.lax easier, because only the functions will be copied over. PiperOrigin-RevId: 430692713 | 24 February 2022, 14:57:56 UTC |
| f51a05a | Peter Hawkins | 17 February 2022, 20:23:44 UTC | Remove jax.ops.index... functions. These functions have been deprecated and have issued a DeprecationWarning since jax 0.2.22 in October 2021. | 24 February 2022, 14:36:28 UTC |
| 3948fde | jax authors | 24 February 2022, 13:37:02 UTC | Merge pull request #9052 from jpuigcerver:main PiperOrigin-RevId: 430680329 | 24 February 2022, 13:37:02 UTC |
| a9a827e | Yash Katariya | 23 February 2022, 23:22:05 UTC | Add .device() method to _DeviceArray PiperOrigin-RevId: 430551695 | 23 February 2022, 23:22:37 UTC |
| 98e114d | Yash Katariya | 23 February 2022, 22:26:28 UTC | Rename `unmapped_local_out_avals` to `out_avals` since it can contain global avals (via GDA) as well as local avals. PiperOrigin-RevId: 430539281 | 23 February 2022, 22:27:00 UTC |
| e2834d8 | Yash Katariya | 23 February 2022, 21:57:59 UTC | Fix the gpu tests that were failing with Future warning PiperOrigin-RevId: 430532523 | 23 February 2022, 21:58:28 UTC |
| c041694 | jax authors | 23 February 2022, 20:00:20 UTC | Merge pull request #8395 from sharadmv:name-stack-mechanism PiperOrigin-RevId: 430506453 | 23 February 2022, 20:00:20 UTC |
| 687a763 | Yash Katariya | 23 February 2022, 18:46:27 UTC | Deprecate `maps.mesh` and replace it with `maps.Mesh`. PiperOrigin-RevId: 430489855 | 23 February 2022, 18:47:06 UTC |
| 1b79caa | Sharad Vikram | 28 October 2021, 18:06:58 UTC | Add separate mechanism for threading name stacks to the lowering | 23 February 2022, 17:59:09 UTC |
| e96b91d | Yash Katariya | 23 February 2022, 17:54:43 UTC | Make the name `pjit` appear in xprof and mhlo module name. Before everything was `xmap` even for `pjit`. PiperOrigin-RevId: 430477182 | 23 February 2022, 17:55:12 UTC |
| 7edd964 | jax authors | 23 February 2022, 16:43:09 UTC | Merge pull request #9663 from mattjj:checkify-scan-debug PiperOrigin-RevId: 430462126 | 23 February 2022, 16:43:09 UTC |
| 45d3ddd | lenamartens | 22 February 2022, 23:08:22 UTC | Fix tests and handle cond consts. | 23 February 2022, 16:11:09 UTC |
| 4b1d0a4 | Matthew Johnson | 22 February 2022, 18:31:05 UTC | fixing scan and other control flow Co-authored-by: Lena Martens <lenamartens@google.com> | 23 February 2022, 15:39:22 UTC |
| ab15db7 | jax authors | 22 February 2022, 23:47:33 UTC | Merge pull request #9613 from mattjj:vmap-binop-batchers-use-singletons PiperOrigin-RevId: 430313560 | 22 February 2022, 23:47:33 UTC |
| 97b1bd3 | jax authors | 22 February 2022, 20:21:44 UTC | Merge pull request #9636 from LenaMartens:changelist/429277776 PiperOrigin-RevId: 430268738 | 22 February 2022, 20:21:44 UTC |
| afdb7b6 | Yash Katariya | 22 February 2022, 18:34:13 UTC | Make local_to_global and global_to_local private. PiperOrigin-RevId: 430243584 | 22 February 2022, 18:34:51 UTC |
| 2641f06 | Adam Paszke | 22 February 2022, 14:01:36 UTC | Always treat all mesh axes controlled by xmap as MANUAL PiperOrigin-RevId: 430192736 | 22 February 2022, 14:02:09 UTC |
| 6423741 | tornikeo | 22 February 2022, 09:49:48 UTC | Add copy button to code snippets in documentation | 22 February 2022, 09:49:48 UTC |
| a65841f | jax authors | 21 February 2022, 08:04:32 UTC | Merge pull request #9617 from froystig:prng-array-tests PiperOrigin-RevId: 429976476 | 21 February 2022, 08:04:32 UTC |
| 660616c | jax authors | 19 February 2022, 21:59:38 UTC | Merge pull request #9607 from jakevdp:scatter-apply PiperOrigin-RevId: 429797642 | 19 February 2022, 21:59:38 UTC |
| 23d5eb0 | Rob Suderman | 19 February 2022, 06:58:51 UTC | IREE's get_default_device_assignment should return List[Device] Previously return List[List[Device]] which is not how the function is used. Updated to use the alternative overload. | 19 February 2022, 07:02:04 UTC |
| 3290dd3 | Yash Katariya | 19 February 2022, 01:27:45 UTC | Make resharding of GDA work if the shape is larger than what it was serialized with. For example: If you serialize with shape (8, 2) and want to deserialize with global shape (12, 2). PiperOrigin-RevId: 429680502 | 19 February 2022, 01:28:13 UTC |
| c161c62 | Yash Katariya | 19 February 2022, 00:23:02 UTC | Finish jax release PiperOrigin-RevId: 429670894 | 19 February 2022, 00:23:39 UTC |
| d1b6f5d | Yash Katariya | 18 February 2022, 22:23:49 UTC | Revert back to adding aval on Device buffers inside local_shards and convert the cached property to just the normal property. This slows down the pjit path because now you are paying the cost to create avals during runtime. PiperOrigin-RevId: 429647845 | 18 February 2022, 22:24:30 UTC |
| 2eeb683 | Lena Martens | 17 February 2022, 12:31:59 UTC | Checkify: address initial feedback. - add a way to run checkify with no errors enabled - clarify "can't be staged" error message - export init_error: a way for users to set a default Error value | 18 February 2022, 19:23:28 UTC |
| 1486be7 | Yash Katariya | 18 February 2022, 19:15:56 UTC | * Make _old_env thread local so that it can be used in multiple threads. * Make nesting for `with mesh` work properly by using a stack. * Allow `Mesh` to be used as a decorator: ``` @pxla.Mesh(mesh_devices, ('x', 'y')) def dec(): return pjit(lambda x: x, in_axis_resources=P('x'), out_axis_resources=None)(x) ``` PiperOrigin-RevId: 429607613 | 18 February 2022, 19:16:23 UTC |
| 8cb1692 | Yuanzhong Xu | 18 February 2022, 18:51:49 UTC | Allow unevenly partitioned sharding_constraints. PiperOrigin-RevId: 429601508 | 18 February 2022, 18:52:20 UTC |
| e13c847 | Jake VanderPlas | 18 February 2022, 17:44:40 UTC | Index update operators: add scatter_apply() | 18 February 2022, 17:44:40 UTC |
| 1baa59c | jax authors | 18 February 2022, 11:23:42 UTC | Merge pull request #9594 from oliverdutton:boolean_jax2tf PiperOrigin-RevId: 429523375 | 18 February 2022, 11:23:42 UTC |
| 607e703 | Jean-Baptiste Lespiau | 18 February 2022, 11:18:19 UTC | Turn execute_replicated into a class so we can access its fields. It's more readable than inspecting the internals of a `functools.partial`. PiperOrigin-RevId: 429523075 | 18 February 2022, 11:18:47 UTC |
| d123a10 | jax authors | 18 February 2022, 09:41:32 UTC | Merge pull request #9631 from gnecula:tf_doc PiperOrigin-RevId: 429508358 | 18 February 2022, 09:41:32 UTC |
| e6d94e1 | oliver | 18 February 2022, 09:10:31 UTC | docs: fix explanation | 18 February 2022, 09:10:31 UTC |
| f926d08 | oliver | 18 February 2022, 09:07:29 UTC | docs: spelling mistake fixed | 18 February 2022, 09:07:29 UTC |
| 3576388 | George Necula | 18 February 2022, 08:38:10 UTC | [jax2tf] Added more links to the documentation | 18 February 2022, 08:38:10 UTC |
| 20635e4 | oliverdutton | 18 February 2022, 08:20:48 UTC | fix: simplify reduce_min with reduce_all | 18 February 2022, 08:20:48 UTC |
| b78ec00 | Yash Katariya | 18 February 2022, 05:06:44 UTC | Fix a typo (remove `== ss`). PiperOrigin-RevId: 429472881 | 18 February 2022, 05:07:14 UTC |
| a430d0f | Yash Katariya | 18 February 2022, 03:27:33 UTC | Show all the device buffer shapes in the error message. PiperOrigin-RevId: 429458927 | 18 February 2022, 03:28:04 UTC |
| d22129f | jax authors | 17 February 2022, 23:15:29 UTC | Merge pull request #9620 from jakevdp:deprecate-jax-test-case PiperOrigin-RevId: 429413723 | 17 February 2022, 23:15:29 UTC |
| da3aaa1 | Jake VanderPlas | 17 February 2022, 22:58:58 UTC | Add deprecation warning to JaxTestCase and JaxTestLoader | 17 February 2022, 22:58:58 UTC |
| 50f4b58 | Rob Suderman | 17 February 2022, 22:31:23 UTC | Fix Iree backend to for copy_to_device and executable results Executable results can be a tuple, if so iterate over the entires. Copy to device should just return the the IREE buffer as device buffer management is still in progress. | 17 February 2022, 22:31:23 UTC |
| e545daa | jax authors | 17 February 2022, 22:23:13 UTC | Merge pull request #9621 from hawkinsp:docs3 PiperOrigin-RevId: 429401742 | 17 February 2022, 22:23:13 UTC |
| 3e5ecfe | Peter Hawkins | 17 February 2022, 21:05:33 UTC | Add jax.distributed and jax.dlpack to the docs. Reorder the doc modules into something closer to alphabetical order. Add missing functions from jax.scipy.linalg and jax.scipy.signal to the docs. | 17 February 2022, 21:10:07 UTC |
| 54a6e4d | jax authors | 17 February 2022, 20:46:12 UTC | Merge pull request #9422 from yotarok:signal_stft PiperOrigin-RevId: 429377655 | 17 February 2022, 20:46:12 UTC |
| 5bb140f | jax authors | 17 February 2022, 20:45:55 UTC | Merge pull request #9615 from froystig:jnp-expand-dims-error PiperOrigin-RevId: 429376858 | 17 February 2022, 20:45:55 UTC |
| 6bb58e6 | Yash Katariya | 17 February 2022, 20:36:44 UTC | Xmap GDA integration. Non-contiguous mesh is allowed! PiperOrigin-RevId: 429376557 | 17 February 2022, 20:41:12 UTC |
| 83a5020 | jax authors | 17 February 2022, 20:36:21 UTC | Merge pull request #9616 from hawkinsp:doc PiperOrigin-RevId: 429376341 | 17 February 2022, 20:36:21 UTC |
| 88c6b84 | Roy Frostig | 17 February 2022, 19:23:15 UTC | in tests, compare jnp operations on PRNGKeyArrays to the same on jnp arrays | 17 February 2022, 19:43:09 UTC |
| d704c15 | Peter Hawkins | 17 February 2022, 19:31:16 UTC | Clarify the NVidia driver version requirements. | 17 February 2022, 19:37:29 UTC |
| 35fab1a | Roy Frostig | 17 February 2022, 19:26:35 UTC | err on repeated axes to `expand_dims`, as numpy does | 17 February 2022, 19:27:20 UTC |
| 032bfe0 | jax authors | 17 February 2022, 18:25:29 UTC | Merge pull request #9609 from froystig:prng-array-stack PiperOrigin-RevId: 429342174 | 17 February 2022, 18:25:29 UTC |
| 6ccbc86 | oliverdutton | 17 February 2022, 16:25:06 UTC | refactor: rename, reorder, rewrite | 17 February 2022, 16:25:06 UTC |
| 73f2370 | Lena Martens | 17 February 2022, 12:30:27 UTC | Checkify: explicitly export public API, hide private symbols. PiperOrigin-RevId: 429277551 | 17 February 2022, 12:30:59 UTC |
| 57f4232 | Adam Paszke | 17 February 2022, 11:27:38 UTC | Fix uninitialized axis_env error when MLIR lowering is disabled PiperOrigin-RevId: 429267926 | 17 February 2022, 11:28:13 UTC |
| 15295a8 | jax authors | 17 February 2022, 11:02:47 UTC | Merge pull request #9544 from SaturdayGenfo:adds-matrix-sqrt PiperOrigin-RevId: 429264231 | 17 February 2022, 11:02:47 UTC |
| b657c23 | oliverdutton | 17 February 2022, 10:02:09 UTC | fix: linting errors | 17 February 2022, 10:02:34 UTC |
| bd2a6a0 | Yash Katariya | 17 February 2022, 07:15:04 UTC | `del self._old_env` so that you can use `with global_mesh` multiple times. This was Matt's idea. Co-authored-by: Matthew Johnson <mattjj@google.com> PiperOrigin-RevId: 429229193 | 17 February 2022, 07:15:36 UTC |
| e0fb424 | Matthew Johnson | 17 February 2022, 07:11:22 UTC | use singleton dims in broadcasting binop batchers | 17 February 2022, 07:11:22 UTC |
