67fad64 | jax authors | 04 May 2021, 00:51:13 UTC | Merge pull request #6632 from skye:version PiperOrigin-RevId: 371816004 | 04 May 2021, 00:51:13 UTC |
1246163 | jax authors | 04 May 2021, 00:27:15 UTC | Merge pull request #6633 from zhangqiaorjc:jaxlib_rel PiperOrigin-RevId: 371812586 | 04 May 2021, 00:27:15 UTC |
24d126b | Qiao Zhang | 04 May 2021, 00:20:15 UTC | Update XLA. | 04 May 2021, 00:20:15 UTC |
c7485b7 | Skye Wanderman-Milne | 03 May 2021, 23:32:00 UTC | Bump jax version and changelog for jax 0.1.13 release | 03 May 2021, 23:32:00 UTC |
e29a9c4 | Qiao Zhang | 03 May 2021, 22:59:00 UTC | [JAX] Bump xla_client._version and gate _ALLOW_ARG_PRUNING on the new version. PiperOrigin-RevId: 371796791 | 03 May 2021, 22:59:25 UTC |
850bd66 | Qiao Zhang | 03 May 2021, 18:40:59 UTC | [JAX] Prune unused inputs in jit. - Python part based on: https://github.com/google/jax/pull/6567 - Added cpp_jit path to handle pruned args PiperOrigin-RevId: 371743277 | 03 May 2021, 18:41:29 UTC |
e6bdcbb | jax authors | 03 May 2021, 17:48:02 UTC | Merge pull request #6625 from hawkinsp:xla PiperOrigin-RevId: 371730573 | 03 May 2021, 17:48:02 UTC |
efe4a1f | Peter Hawkins | 03 May 2021, 17:38:12 UTC | Update XLA. | 03 May 2021, 17:38:12 UTC |
ffecef5 | jax authors | 03 May 2021, 17:26:33 UTC | Merge pull request #6611 from njunge94:pip-install-fork PiperOrigin-RevId: 371724836 | 03 May 2021, 17:26:33 UTC |
321ecbd | jax authors | 03 May 2021, 17:05:55 UTC | Merge pull request #6620 from hawkinsp:cuda PiperOrigin-RevId: 371720280 | 03 May 2021, 17:05:55 UTC |
2c92bc9 | Peter Hawkins | 03 May 2021, 14:48:18 UTC | [JAX] Include pre-transformed stack traces as additional context to JAX exceptions, where present. PiperOrigin-RevId: 371695248 | 03 May 2021, 14:48:47 UTC |
f9899df | Peter Hawkins | 03 May 2021, 14:30:14 UTC | Readd CUDA 10.1 to the set of released CUDA wheels. Users are still depending on CUDA 10.1. | 03 May 2021, 14:30:14 UTC |
75b00a1 | jax authors | 02 May 2021, 05:18:13 UTC | Copybara import of the project: -- 3c400a3e588abf9e2259119c50343cba6f3477f1 by Matthew Johnson <mattjj@google.com>: add 'inline' option to xla_call for jaxpr inlining -- fe297e39ca37896b75d7943b9b77c0b53fad13ee by Matthew Johnson <mattjj@google.com>: add 'inline' to jit docstring -- ff6866c4b3757cde66fe659c2f27d8aeff024e8f by Matthew Johnson <mattjj@google.com>: new_sublevel in jax2tf PiperOrigin-RevId: 371542778 | 02 May 2021, 05:18:39 UTC |
e8f209c | jax authors | 02 May 2021, 04:53:34 UTC | Merge pull request #6584 from google:jit-inline-2 PiperOrigin-RevId: 371541392 | 02 May 2021, 04:53:34 UTC |
712d5f7 | jax authors | 02 May 2021, 03:19:41 UTC | Merge pull request #6618 from imh:autodidax_broadcasting_error PiperOrigin-RevId: 371535729 | 02 May 2021, 03:19:41 UTC |
74b6762 | Ian Horn | 02 May 2021, 02:07:58 UTC | Fixed error in autodidax broadasting batching | 02 May 2021, 02:07:58 UTC |
89a6b6d | Peter Hawkins | 01 May 2021, 23:12:44 UTC | Fix stack trace filtering for paths that do not exist. PiperOrigin-RevId: 371521548 | 01 May 2021, 23:13:07 UTC |
a80aab0 | Peter Hawkins | 01 May 2021, 21:57:15 UTC | Fix cloudpickle breakage. PiperOrigin-RevId: 371517484 | 01 May 2021, 21:57:41 UTC |
bcd33b7 | jax authors | 01 May 2021, 21:26:47 UTC | Merge pull request #6617 from hawkinsp:xla PiperOrigin-RevId: 371515560 | 01 May 2021, 21:26:47 UTC |
cacc5af | Peter Hawkins | 01 May 2021, 21:05:14 UTC | Update XLA. | 01 May 2021, 21:05:14 UTC |
ff6866c | Matthew Johnson | 01 May 2021, 19:41:41 UTC | new_sublevel in jax2tf | 01 May 2021, 19:41:41 UTC |
e8c3406 | Peter Hawkins | 01 May 2021, 19:41:10 UTC | [JAX] Switch the order of the filtered and unfiltered stack traces in exceptions. After this change, the filtered stack trace is attached to the main exception, and the unfiltered stack trace becomes a __cause__ exception. PiperOrigin-RevId: 371509766 | 01 May 2021, 19:41:37 UTC |
fe297e3 | Matthew Johnson | 01 May 2021, 19:32:44 UTC | add 'inline' to jit docstring | 01 May 2021, 19:32:44 UTC |
f559a9b | Nicholas Junge | 01 May 2021, 17:32:19 UTC | Fix pip install instruction for forks in contributing section | 01 May 2021, 17:32:19 UTC |
6ce4ef4 | jax authors | 01 May 2021, 00:12:02 UTC | Merge pull request #6593 from njunge94:axis_concatenation PiperOrigin-RevId: 371431622 | 01 May 2021, 00:12:02 UTC |
71c7f29 | jax authors | 30 April 2021, 23:31:38 UTC | Merge pull request #6606 from jakevdp:numpy-version PiperOrigin-RevId: 371425726 | 30 April 2021, 23:31:38 UTC |
e985e86 | Nicholas Junge | 30 April 2021, 23:05:22 UTC | Implement jnp.r_ and jnp.c_ | 30 April 2021, 23:05:22 UTC |
7d61def | Jake VanderPlas | 30 April 2021, 20:56:00 UTC | Cleanup: use numpy_version rather than np.__version__ | 30 April 2021, 20:56:00 UTC |
9c63c77 | jax authors | 30 April 2021, 18:33:07 UTC | Merge pull request #6597 from lgeiger:np.broadcast_shapes PiperOrigin-RevId: 371371698 | 30 April 2021, 18:33:07 UTC |
a613ce1 | Lukas Geiger | 30 April 2021, 17:32:51 UTC | Export `jnp.broadcast_shapes` as user facing function | 30 April 2021, 17:32:51 UTC |
893b5a0 | Adam Paszke | 30 April 2021, 16:56:53 UTC | Make pjit into an initial style primitive Partly to make it more robust (e.g. we no longer need to implement post_process_call), partly beacuse it is not really a call primitive (it modifies the argument and return avals in a multiprocess mesh), and partly as an experiment to see how difficult would it be to actually make it more autodidax-like. Overall, it seems like a mixed bag, maybe slightly positive. The thunks are gone which is nice, but one has to be much more careful when dealing with avals. PiperOrigin-RevId: 371352737 | 30 April 2021, 16:57:36 UTC |
c92a9c0 | jax authors | 30 April 2021, 01:35:59 UTC | Merge pull request #6596 from lgeiger:simplify-numpy-wraps PiperOrigin-RevId: 371247269 | 30 April 2021, 01:35:59 UTC |
ca35927 | Lukas Geiger | 29 April 2021, 22:57:57 UTC | Remove docs fallback for numpy < 1.15 | 29 April 2021, 22:57:57 UTC |
bd35379 | jax authors | 29 April 2021, 21:02:56 UTC | Merge pull request #6591 from jakevdp:tpu-test-setup PiperOrigin-RevId: 371198936 | 29 April 2021, 21:02:56 UTC |
ccc423d | Peter Hawkins | 29 April 2021, 20:36:38 UTC | [JAX] Enable filtered tracebacks on Python 3.6. [XLA:Python] Add support for converting a fast-traceback into a Python exception traceback. Add a helper for building traceback objects on Python 3.6. On Python 3.7+ this can be done by calling the traceback type, and we can in essence backport that implementation to Python 3.6. Consolidate the py_traceback and traceback modules. PiperOrigin-RevId: 371193212 | 29 April 2021, 20:37:21 UTC |
7c9357e | jax authors | 29 April 2021, 20:33:19 UTC | Merge pull request #6588 from google:nn-tanh PiperOrigin-RevId: 371192342 | 29 April 2021, 20:33:19 UTC |
ef91c0d | jax authors | 29 April 2021, 19:30:19 UTC | Merge pull request #6590 from jakevdp:nan-quantile PiperOrigin-RevId: 371178272 | 29 April 2021, 19:30:19 UTC |
75894b9 | jax authors | 29 April 2021, 19:03:58 UTC | Merge pull request #6589 from jakevdp:lax-doc PiperOrigin-RevId: 371172471 | 29 April 2021, 19:03:58 UTC |
e1f2c04 | Jake VanderPlas | 29 April 2021, 18:13:52 UTC | test: use colab_tpu.setup_tpu() in TPU test notebook | 29 April 2021, 18:13:52 UTC |
c9d6def | Jake VanderPlas | 29 April 2021, 18:00:22 UTC | BUG: ensure jnp.quantile() returns NaN for arrays containing NaN | 29 April 2021, 18:00:22 UTC |
6fd806b | jax authors | 29 April 2021, 17:56:07 UTC | Merge pull request #6585 from google:tree-map-alias PiperOrigin-RevId: 371156638 | 29 April 2021, 17:56:07 UTC |
6826021 | jax authors | 29 April 2021, 17:35:41 UTC | Merge pull request #6581 from hawkinsp:libdevice PiperOrigin-RevId: 371151412 | 29 April 2021, 17:35:41 UTC |
71a25cd | Jake VanderPlas | 29 April 2021, 16:48:52 UTC | DOC: add examples to lax function docstrings | 29 April 2021, 16:48:52 UTC |
e968672 | Matthew Johnson | 29 April 2021, 15:25:26 UTC | add tanh to jax.nn package | 29 April 2021, 15:26:38 UTC |
c983d3c | Peter Hawkins | 28 April 2021, 15:43:50 UTC | Bundle libdevice.10.bc with jaxlib wheels. libdevice.10.bc is a redistributable part of the CUDA SDK. This avoids problems trying to locate a copy of libdevice inside the user's CUDA installation. | 29 April 2021, 14:26:03 UTC |
23cbcbe | jax authors | 29 April 2021, 09:11:09 UTC | Merge pull request #6392 from apaszke:axis-name-top-trace PiperOrigin-RevId: 371075561 | 29 April 2021, 09:11:09 UTC |
b8f9dd6 | Roy Frostig | 29 April 2021, 02:44:20 UTC | unify tree_map and tree_multimap | 29 April 2021, 02:59:31 UTC |
3c400a3 | Matthew Johnson | 23 April 2021, 03:56:17 UTC | add 'inline' option to xla_call for jaxpr inlining | 29 April 2021, 02:38:15 UTC |
4fcfaeb | jax authors | 28 April 2021, 13:04:42 UTC | Merge pull request #6577 from gnecula:hcb1 PiperOrigin-RevId: 370886452 | 28 April 2021, 13:04:42 UTC |
8df502a | Adam Paszke | 09 April 2021, 12:43:40 UTC | Use the axis names attached to a primitive when selecting the top trace This is useful e.g. for handling psums of values that are not sharded, but are also not statically known constants that we can fold. | 28 April 2021, 09:46:24 UTC |
d762ec1 | George Necula | 28 April 2021, 09:22:32 UTC | [host_callback] Minor fix to use the new xla_shape.is_token | 28 April 2021, 09:22:32 UTC |
2c7556e | jax authors | 28 April 2021, 03:50:13 UTC | Merge pull request #6574 from jakevdp:gpu-test PiperOrigin-RevId: 370820776 | 28 April 2021, 03:50:13 UTC |
69e9c64 | jax authors | 28 April 2021, 01:52:40 UTC | Merge pull request #6575 from hawkinsp:cusparse PiperOrigin-RevId: 370808176 | 28 April 2021, 01:52:40 UTC |
a0c96b5 | Peter Hawkins | 28 April 2021, 00:28:30 UTC | Fix cusparse kernel build. The build_wheel.py script was copying the wrong module. In addition the CUDA stubs from the TF repo were missing a number of cusparse symbols. The updated TF includes the correct stubs. | 28 April 2021, 00:28:30 UTC |
e3ac7c2 | Jake VanderPlas | 27 April 2021, 22:42:39 UTC | sparse: make gpu translation test more clear | 27 April 2021, 22:42:39 UTC |
4e49089 | jax authors | 27 April 2021, 22:36:54 UTC | Merge pull request #6573 from jakevdp:contributing PiperOrigin-RevId: 370777338 | 27 April 2021, 22:36:54 UTC |
ba83ab5 | Jake VanderPlas | 27 April 2021, 22:00:21 UTC | Add expanded contributing documentation | 27 April 2021, 22:00:21 UTC |
342d624 | George Necula | 27 April 2021, 17:29:39 UTC | [host_callback] Add support for pjit of host_callback. Currently, all XLA side-effect ops inside a sharded computation must have explicit sharding. This includes the outfeed and infeed used by host_callback. The implementation here uses AssignDevice sharding for both the outfeed and the infeed. This means that before outfeed, the devices will do an all_gather and the first device will make the outfeed. The host callback will receive a single outfeed with the entire array, and is supposed to return the entire array. This gets sent to the same device that issued to outfeed, which is responsible to send the respective slices to the other participating devices. PiperOrigin-RevId: 370711606 | 27 April 2021, 17:30:13 UTC |
aa9d350 | Adam Paszke | 27 April 2021, 14:11:41 UTC | Fix an assertion in xmap code This has gone under the radar because it's valid in tests and it's in code sufficiently new that I don't expect anyone to be using it just yet. PiperOrigin-RevId: 370672928 | 27 April 2021, 14:12:23 UTC |
da0adb2 | jax authors | 27 April 2021, 13:18:39 UTC | Merge pull request #6572 from gnecula:tf_call PiperOrigin-RevId: 370664497 | 27 April 2021, 13:18:39 UTC |
80cd273 | George Necula | 27 April 2021, 12:21:20 UTC | [jax2tf] Improve the handling of variable capture for jax2tf.call Fixes: #6050 | 27 April 2021, 12:51:04 UTC |
9f2ac6e | Adam Paszke | 27 April 2021, 09:19:18 UTC | Add resource type-checking rules for pjit and with_sharding_constraint To prevent people from doubly-sharding values that are sharded by outer xmaps. PiperOrigin-RevId: 370635057 | 27 April 2021, 09:19:43 UTC |
43d2733 | jax authors | 27 April 2021, 08:28:07 UTC | Merge pull request #6377 from boyentenbi:changelist/367423156 PiperOrigin-RevId: 370628969 | 27 April 2021, 08:28:07 UTC |
c29400e | jax authors | 27 April 2021, 01:59:26 UTC | Merge pull request #6563 from hawkinsp:cuda PiperOrigin-RevId: 370589037 | 27 April 2021, 01:59:26 UTC |
07ae0c7 | jax authors | 27 April 2021, 00:48:57 UTC | Merge pull request #6466 from jakevdp:sparse-op-objects PiperOrigin-RevId: 370579960 | 27 April 2021, 00:48:57 UTC |
b16b587 | Jake VanderPlas | 27 April 2021, 00:40:17 UTC | sparse: add initial higher-level sparse objects | 27 April 2021, 00:40:17 UTC |
d6c1aad | jax authors | 26 April 2021, 20:52:56 UTC | Merge pull request #6550 from jakevdp:dynamic-slice-doc PiperOrigin-RevId: 370534106 | 26 April 2021, 20:52:56 UTC |
d060646 | Adam Paszke | 26 April 2021, 18:41:26 UTC | Fix the batching rule for named reductions PiperOrigin-RevId: 370505998 | 26 April 2021, 18:41:58 UTC |
a0c2c33 | jax authors | 26 April 2021, 17:50:11 UTC | Merge pull request #6565 from jakevdp:cuda-spec PiperOrigin-RevId: 370493528 | 26 April 2021, 17:50:11 UTC |
bdcce8f | Jake VanderPlas | 26 April 2021, 16:22:52 UTC | setup.py: make cuda extras specifier more consistent | 26 April 2021, 16:22:52 UTC |
64130da | Adam Paszke | 26 April 2021, 14:51:19 UTC | Fix xmap nesting in SPMD lowering mode Previous implementation has been incorrectly accumulating the SPMD axis partitioning into an existing dimension, while what we really need is to create a separate dimension for every annotation. PiperOrigin-RevId: 370455989 | 26 April 2021, 14:51:44 UTC |
79a7f7b | Peter Hawkins | 26 April 2021, 13:38:31 UTC | Don't build CUDA 11.2 wheels. Update XLA. CUDA 11.1 wheels are compatible with CUDA versions 11.1+, since NVidia now promises enhanced version compatibility between CUDA minor releases starting with CUDA 11.1 | 26 April 2021, 13:43:29 UTC |
65137ab | Adam Paszke | 26 April 2021, 13:41:44 UTC | Fix two issues caused by the recent pjit parsing patch 1. We were accidentally passing in ParsedPartitionSpecs as user inputs in the AD rules for `with_sharding_constraint`. 2. The `_check_shapes_against_resources` function assumed it was getting called with avals, but was really supplied real inputs and outputs. Since those include raw Python number objects, one could easily get a failure when we queried there shape. PiperOrigin-RevId: 370445989 | 26 April 2021, 13:42:13 UTC |
6cbab8e | Adam Paszke | 26 April 2021, 10:45:31 UTC | Parse and verify pjit in/out_axis_resources This lets us focus the handling of all the quirks of `PartitionSpec` in one place and should make it easier to extend support to xmap-like specs in the future. This also adds new checks for repeated assignment of mesh axes which would previously went unnoticed (the first annotation would simply get ignored). PiperOrigin-RevId: 370425384 | 26 April 2021, 10:45:56 UTC |
a9986ce | jax authors | 25 April 2021, 03:21:55 UTC | Merge pull request #6557 from google:enable-working-test PiperOrigin-RevId: 370294198 | 25 April 2021, 03:21:55 UTC |
8f43453 | Matthew Johnson | 24 April 2021, 22:18:26 UTC | re-enable a working test | 24 April 2021, 22:18:26 UTC |
887ed30 | jax authors | 23 April 2021, 19:29:49 UTC | Merge pull request #6554 from hawkinsp:gmres PiperOrigin-RevId: 370134381 | 23 April 2021, 19:29:49 UTC |
fb74246 | Peter Hawkins | 23 April 2021, 19:13:40 UTC | Disable gmres test on CPU. This test has started failing at LLVM head; disabling it while we debug. | 23 April 2021, 19:13:40 UTC |
b83931a | jax authors | 23 April 2021, 18:15:49 UTC | Merge pull request #6552 from jakevdp:fix-sparse-nnz PiperOrigin-RevId: 370118143 | 23 April 2021, 18:15:49 UTC |
c9c3cb6 | jax authors | 23 April 2021, 17:47:17 UTC | Merge pull request #6534 from skye:test_warnings PiperOrigin-RevId: 370111662 | 23 April 2021, 17:47:17 UTC |
f6d9e5a | Jake VanderPlas | 23 April 2021, 17:36:28 UTC | [sparse] fix conversion issue when nnz is larger than necessary | 23 April 2021, 17:36:28 UTC |
3ff51bb | Skye Wanderman-Milne | 21 April 2021, 21:31:57 UTC | Use pytest's filterwarnings feature instead of filtering each test case. We often forget to put the per-test-case decorators, resulting in test failures in cases not covered by github CI (e.g. Cloud TPU tests). This change filters the "experimental feature" warnings by default. | 23 April 2021, 17:28:22 UTC |
6969062 | jax authors | 23 April 2021, 17:07:19 UTC | Merge pull request #6549 from google:lower-fori-to-scan-2 PiperOrigin-RevId: 370102373 | 23 April 2021, 17:07:19 UTC |
f0264ea | jax authors | 23 April 2021, 17:03:36 UTC | Merge pull request #6548 from LenaMartens:changelist/370070595 PiperOrigin-RevId: 370102278 | 23 April 2021, 17:03:36 UTC |
ca684df | Jake VanderPlas | 23 April 2021, 16:10:43 UTC | DOC: add example for lax.dynamic_update_slice | 23 April 2021, 16:10:43 UTC |
4ac8937 | Matthew Johnson | 23 April 2021, 15:31:11 UTC | enable fori-to-scan lowering | 23 April 2021, 15:49:13 UTC |
b244e2b | Lena Martens | 23 April 2021, 13:43:20 UTC | Add eval_shape to the UnexpectedTracerError too. | 23 April 2021, 13:46:34 UTC |
973ca07 | Adam Paszke | 23 April 2021, 10:42:14 UTC | Add stricter resource overlap checking We've had some checks for coinciding logical axes mapped to the same resources in the existing xmap code, but they were quite lax. This adds a proper type checker and a bunch of tests to verify that we can catch the interesting failure cases. PiperOrigin-RevId: 370051512 | 23 April 2021, 10:42:34 UTC |
615f9ad | jax authors | 23 April 2021, 00:28:33 UTC | Merge pull request #6536 from pschuh:compile-replicated PiperOrigin-RevId: 369985107 | 23 April 2021, 00:28:33 UTC |
23f847e | Adam Paszke | 22 April 2021, 22:30:03 UTC | Make the initial-style -> final-style conversion rule based Also, add a rule for pjit to make sure that we can eval jaxprs that contain pjits. PiperOrigin-RevId: 369964136 | 22 April 2021, 22:30:30 UTC |
594309c | jax authors | 22 April 2021, 17:08:42 UTC | Merge pull request #6537 from gnecula:disable_tests PiperOrigin-RevId: 369895220 | 22 April 2021, 17:08:42 UTC |
9d22a0a | jax authors | 22 April 2021, 16:15:19 UTC | Merge pull request #6467 from lgeiger:conv-preferred-element-type PiperOrigin-RevId: 369884941 | 22 April 2021, 16:15:19 UTC |
eb9d6e4 | Peter Choy | 08 April 2021, 14:08:51 UTC | Pass axis name to _match_axes and add to error message. | 22 April 2021, 13:34:04 UTC |
454f5e6 | Adam Paszke | 22 April 2021, 09:46:42 UTC | Enable nesting xmaps in SPMD lowering mode This uses the recently added ability to modify the `BatchTrace` to add a new `SPMDBatchTrace`, which additionally fills in `spmd_in_axes` and `spmd_out_axes` of xmap primitives. Those fields are necessary, because XLA does not allow us to emit partial `OpSharding` annotations, meaning that we have to track where the positional axes of outer xmaps are inserted at the boundaries of inner xmaps. Otherwise, XLA could misinterpret our intention and unnecessarily force replication along the mesh axes used by the outer xmaps. PiperOrigin-RevId: 369831571 | 22 April 2021, 09:47:07 UTC |
f7f4269 | Lukas Geiger | 21 April 2021, 21:58:34 UTC | Add support for `preferred_element_type` arg in convolutions | 22 April 2021, 08:29:31 UTC |
6ed0686 | George Necula | 22 April 2021, 05:14:56 UTC | Extended the comment about disabling the test_multiple_barriers | 22 April 2021, 05:14:56 UTC |
6e5880a | jax authors | 21 April 2021, 23:23:23 UTC | Merge pull request #6532 from jakevdp:unique-axis PiperOrigin-RevId: 369758257 | 21 April 2021, 23:23:23 UTC |
bb543f2 | Jake VanderPlas | 21 April 2021, 23:00:14 UTC | jnp.unique: add support for axis argument | 21 April 2021, 23:00:14 UTC |
bb96ae7 | Parker Schuh | 21 April 2021, 22:32:47 UTC | Change `wrap_execute_replicated` -> `compile_replicated`. This avoids accessing the compiled object in order to pass devices to shard_args. This allows the compiled object to be constructed lazily inside the backend specific `compile_replicated`. | 21 April 2021, 22:32:47 UTC |
90d606f | Jake VanderPlas | 21 April 2021, 21:51:46 UTC | Remove jax.experimental.doubledouble PiperOrigin-RevId: 369740697 | 21 April 2021, 21:52:14 UTC |