swh:1:snp:71afc25eb6e6e055a37a962e6b91010ec35e397f

sort by:
Revision Author Date Message Commit Date
67fad64 Merge pull request #6632 from skye:version PiperOrigin-RevId: 371816004 04 May 2021, 00:51:13 UTC
1246163 Merge pull request #6633 from zhangqiaorjc:jaxlib_rel PiperOrigin-RevId: 371812586 04 May 2021, 00:27:15 UTC
24d126b Update XLA. 04 May 2021, 00:20:15 UTC
c7485b7 Bump jax version and changelog for jax 0.1.13 release 03 May 2021, 23:32:00 UTC
e29a9c4 [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 [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 Merge pull request #6625 from hawkinsp:xla PiperOrigin-RevId: 371730573 03 May 2021, 17:48:02 UTC
efe4a1f Update XLA. 03 May 2021, 17:38:12 UTC
ffecef5 Merge pull request #6611 from njunge94:pip-install-fork PiperOrigin-RevId: 371724836 03 May 2021, 17:26:33 UTC
321ecbd Merge pull request #6620 from hawkinsp:cuda PiperOrigin-RevId: 371720280 03 May 2021, 17:05:55 UTC
2c92bc9 [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 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 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 Merge pull request #6584 from google:jit-inline-2 PiperOrigin-RevId: 371541392 02 May 2021, 04:53:34 UTC
712d5f7 Merge pull request #6618 from imh:autodidax_broadcasting_error PiperOrigin-RevId: 371535729 02 May 2021, 03:19:41 UTC
74b6762 Fixed error in autodidax broadasting batching 02 May 2021, 02:07:58 UTC
89a6b6d Fix stack trace filtering for paths that do not exist. PiperOrigin-RevId: 371521548 01 May 2021, 23:13:07 UTC
a80aab0 Fix cloudpickle breakage. PiperOrigin-RevId: 371517484 01 May 2021, 21:57:41 UTC
bcd33b7 Merge pull request #6617 from hawkinsp:xla PiperOrigin-RevId: 371515560 01 May 2021, 21:26:47 UTC
cacc5af Update XLA. 01 May 2021, 21:05:14 UTC
ff6866c new_sublevel in jax2tf 01 May 2021, 19:41:41 UTC
e8c3406 [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 add 'inline' to jit docstring 01 May 2021, 19:32:44 UTC
f559a9b Fix pip install instruction for forks in contributing section 01 May 2021, 17:32:19 UTC
6ce4ef4 Merge pull request #6593 from njunge94:axis_concatenation PiperOrigin-RevId: 371431622 01 May 2021, 00:12:02 UTC
71c7f29 Merge pull request #6606 from jakevdp:numpy-version PiperOrigin-RevId: 371425726 30 April 2021, 23:31:38 UTC
e985e86 Implement jnp.r_ and jnp.c_ 30 April 2021, 23:05:22 UTC
7d61def Cleanup: use numpy_version rather than np.__version__ 30 April 2021, 20:56:00 UTC
9c63c77 Merge pull request #6597 from lgeiger:np.broadcast_shapes PiperOrigin-RevId: 371371698 30 April 2021, 18:33:07 UTC
a613ce1 Export `jnp.broadcast_shapes` as user facing function 30 April 2021, 17:32:51 UTC
893b5a0 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 Merge pull request #6596 from lgeiger:simplify-numpy-wraps PiperOrigin-RevId: 371247269 30 April 2021, 01:35:59 UTC
ca35927 Remove docs fallback for numpy < 1.15 29 April 2021, 22:57:57 UTC
bd35379 Merge pull request #6591 from jakevdp:tpu-test-setup PiperOrigin-RevId: 371198936 29 April 2021, 21:02:56 UTC
ccc423d [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 Merge pull request #6588 from google:nn-tanh PiperOrigin-RevId: 371192342 29 April 2021, 20:33:19 UTC
ef91c0d Merge pull request #6590 from jakevdp:nan-quantile PiperOrigin-RevId: 371178272 29 April 2021, 19:30:19 UTC
75894b9 Merge pull request #6589 from jakevdp:lax-doc PiperOrigin-RevId: 371172471 29 April 2021, 19:03:58 UTC
e1f2c04 test: use colab_tpu.setup_tpu() in TPU test notebook 29 April 2021, 18:13:52 UTC
c9d6def BUG: ensure jnp.quantile() returns NaN for arrays containing NaN 29 April 2021, 18:00:22 UTC
6fd806b Merge pull request #6585 from google:tree-map-alias PiperOrigin-RevId: 371156638 29 April 2021, 17:56:07 UTC
6826021 Merge pull request #6581 from hawkinsp:libdevice PiperOrigin-RevId: 371151412 29 April 2021, 17:35:41 UTC
71a25cd DOC: add examples to lax function docstrings 29 April 2021, 16:48:52 UTC
e968672 add tanh to jax.nn package 29 April 2021, 15:26:38 UTC
c983d3c 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 Merge pull request #6392 from apaszke:axis-name-top-trace PiperOrigin-RevId: 371075561 29 April 2021, 09:11:09 UTC
b8f9dd6 unify tree_map and tree_multimap 29 April 2021, 02:59:31 UTC
3c400a3 add 'inline' option to xla_call for jaxpr inlining 29 April 2021, 02:38:15 UTC
4fcfaeb Merge pull request #6577 from gnecula:hcb1 PiperOrigin-RevId: 370886452 28 April 2021, 13:04:42 UTC
8df502a 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 [host_callback] Minor fix to use the new xla_shape.is_token 28 April 2021, 09:22:32 UTC
2c7556e Merge pull request #6574 from jakevdp:gpu-test PiperOrigin-RevId: 370820776 28 April 2021, 03:50:13 UTC
69e9c64 Merge pull request #6575 from hawkinsp:cusparse PiperOrigin-RevId: 370808176 28 April 2021, 01:52:40 UTC
a0c96b5 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 sparse: make gpu translation test more clear 27 April 2021, 22:42:39 UTC
4e49089 Merge pull request #6573 from jakevdp:contributing PiperOrigin-RevId: 370777338 27 April 2021, 22:36:54 UTC
ba83ab5 Add expanded contributing documentation 27 April 2021, 22:00:21 UTC
342d624 [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 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 Merge pull request #6572 from gnecula:tf_call PiperOrigin-RevId: 370664497 27 April 2021, 13:18:39 UTC
80cd273 [jax2tf] Improve the handling of variable capture for jax2tf.call Fixes: #6050 27 April 2021, 12:51:04 UTC
9f2ac6e 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 Merge pull request #6377 from boyentenbi:changelist/367423156 PiperOrigin-RevId: 370628969 27 April 2021, 08:28:07 UTC
c29400e Merge pull request #6563 from hawkinsp:cuda PiperOrigin-RevId: 370589037 27 April 2021, 01:59:26 UTC
07ae0c7 Merge pull request #6466 from jakevdp:sparse-op-objects PiperOrigin-RevId: 370579960 27 April 2021, 00:48:57 UTC
b16b587 sparse: add initial higher-level sparse objects 27 April 2021, 00:40:17 UTC
d6c1aad Merge pull request #6550 from jakevdp:dynamic-slice-doc PiperOrigin-RevId: 370534106 26 April 2021, 20:52:56 UTC
d060646 Fix the batching rule for named reductions PiperOrigin-RevId: 370505998 26 April 2021, 18:41:58 UTC
a0c2c33 Merge pull request #6565 from jakevdp:cuda-spec PiperOrigin-RevId: 370493528 26 April 2021, 17:50:11 UTC
bdcce8f setup.py: make cuda extras specifier more consistent 26 April 2021, 16:22:52 UTC
64130da 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 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 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 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 Merge pull request #6557 from google:enable-working-test PiperOrigin-RevId: 370294198 25 April 2021, 03:21:55 UTC
8f43453 re-enable a working test 24 April 2021, 22:18:26 UTC
887ed30 Merge pull request #6554 from hawkinsp:gmres PiperOrigin-RevId: 370134381 23 April 2021, 19:29:49 UTC
fb74246 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 Merge pull request #6552 from jakevdp:fix-sparse-nnz PiperOrigin-RevId: 370118143 23 April 2021, 18:15:49 UTC
c9c3cb6 Merge pull request #6534 from skye:test_warnings PiperOrigin-RevId: 370111662 23 April 2021, 17:47:17 UTC
f6d9e5a [sparse] fix conversion issue when nnz is larger than necessary 23 April 2021, 17:36:28 UTC
3ff51bb 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 Merge pull request #6549 from google:lower-fori-to-scan-2 PiperOrigin-RevId: 370102373 23 April 2021, 17:07:19 UTC
f0264ea Merge pull request #6548 from LenaMartens:changelist/370070595 PiperOrigin-RevId: 370102278 23 April 2021, 17:03:36 UTC
ca684df DOC: add example for lax.dynamic_update_slice 23 April 2021, 16:10:43 UTC
4ac8937 enable fori-to-scan lowering 23 April 2021, 15:49:13 UTC
b244e2b Add eval_shape to the UnexpectedTracerError too. 23 April 2021, 13:46:34 UTC
973ca07 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 Merge pull request #6536 from pschuh:compile-replicated PiperOrigin-RevId: 369985107 23 April 2021, 00:28:33 UTC
23f847e 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 Merge pull request #6537 from gnecula:disable_tests PiperOrigin-RevId: 369895220 22 April 2021, 17:08:42 UTC
9d22a0a Merge pull request #6467 from lgeiger:conv-preferred-element-type PiperOrigin-RevId: 369884941 22 April 2021, 16:15:19 UTC
eb9d6e4 Pass axis name to _match_axes and add to error message. 22 April 2021, 13:34:04 UTC
454f5e6 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 Add support for `preferred_element_type` arg in convolutions 22 April 2021, 08:29:31 UTC
6ed0686 Extended the comment about disabling the test_multiple_barriers 22 April 2021, 05:14:56 UTC
6e5880a Merge pull request #6532 from jakevdp:unique-axis PiperOrigin-RevId: 369758257 21 April 2021, 23:23:23 UTC
bb543f2 jnp.unique: add support for axis argument 21 April 2021, 23:00:14 UTC
bb96ae7 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 Remove jax.experimental.doubledouble PiperOrigin-RevId: 369740697 21 April 2021, 21:52:14 UTC
back to top