swh:1:snp:71afc25eb6e6e055a37a962e6b91010ec35e397f

sort by:
Revision Author Date Message Commit Date
49b0120 [adx3] Add string labels for second-order primitive subjaxprs. Subjaxpr labels are helpful for understanding their purpose. 16 February 2024, 15:42:10 UTC
03936b5 Add a tuple type 16 February 2024, 15:41:16 UTC
c026ec6 Add a polymorphic list type 12 February 2024, 14:39:08 UTC
e1e9c6e Add cond to adx3 01 February 2024, 15:26:01 UTC
f87c9d0 Factor adx3 to emphasise extending with new types and primitives (but not transformations) 01 February 2024, 03:16:45 UTC
d63c585 Added transposition and linear jaxprs to autodidax3 27 January 2024, 22:06:42 UTC
2b83731 Playing around with a purely "initial style" version of autodidax. 27 January 2024, 03:27:58 UTC
0f760ee Avoid using lambda as the reducer fn Lambdas are represented by their ids in the metadata of lowered HLO (see example below) and they change every time. This makes the compilation cache less effective as it causes the computation's fingerprint to change every time. ``` get-tuple-element.41724 = bf16[8]{0} get-tuple-element(reduce.41723), index=0, metadata={op_name="pjit(_wrapped_fn)/jit(main)/.../reduce[computation=<function _compute_argminmax.<locals>.reducer_fn at 0x7fa6ecfb2200> dimensions=(1,)]" source_file="..." source_line=...} ``` PiperOrigin-RevId: 601910715 27 January 2024, 01:43:57 UTC
ccfe9c1 Merge pull request #19540 from mattjj:remove-hypothesis-test-dependence PiperOrigin-RevId: 601908297 27 January 2024, 01:28:26 UTC
c42305a Merge pull request #19536 from jakevdp:key-reuse-cond PiperOrigin-RevId: 601900128 27 January 2024, 00:43:44 UTC
54d7d5c make hypothesis dependence optional 27 January 2024, 00:31:01 UTC
17935af [key reuse] fix key reuse type for cond with sources 26 January 2024, 22:42:55 UTC
2478f31 remove key array's isinstance-overriding metaclass We don't need to support `isinstance(..., PRNGKeyArray)` on tracers any longer, since `PRNGKeyArray` is no longer a public symbol. PiperOrigin-RevId: 601815616 26 January 2024, 19:16:56 UTC
592809c Roll-back 1ae054b003088d873902fa62cfa8099260471e16 to re-enable nextafter tests Reverts 1ae054b003088d873902fa62cfa8099260471e16 PiperOrigin-RevId: 601814205 26 January 2024, 19:08:43 UTC
cc5f565 Ported a subset of binary operations to lower directly to Triton IR PiperOrigin-RevId: 601806008 26 January 2024, 18:57:01 UTC
8c050ac Merge pull request #19517 from ppham27:changelist/601457375 PiperOrigin-RevId: 601804604 26 January 2024, 18:49:00 UTC
d0008fb Merge pull request #19511 from jakevdp:fix-asarray PiperOrigin-RevId: 601803214 26 January 2024, 18:40:59 UTC
269ad9f Merge pull request #19504 from jakevdp:full-like-device PiperOrigin-RevId: 601803117 26 January 2024, 18:32:53 UTC
cb7a32a Fixed a bug in _reduction_lowering The block argument of tt.reduce is always parameterized by scalars. Note that this bug had no effect on the emitted Triton IR, because the lowering code does not currently rely on avals. PiperOrigin-RevId: 601801294 26 January 2024, 18:24:08 UTC
273cb27 compat.tensor __*__ methods no longer do implicit broadcasting This change makes it simpler to lower binary operations to Triton IR bypassing Triton Python bindings. PiperOrigin-RevId: 601796719 26 January 2024, 18:13:51 UTC
2a8ce9a Merge pull request #19518 from jakevdp:softmax PiperOrigin-RevId: 601796039 26 January 2024, 18:04:48 UTC
9549c74 jnp.full_like & co: support device parameter 26 January 2024, 18:01:54 UTC
d989f50 lax.asarray: avoid explicit device_put 26 January 2024, 17:46:09 UTC
a282d58 nn.softmax: use double-where when where is specified 26 January 2024, 17:45:31 UTC
1ae054b Temporarily disable flaky nextafter tests These are currently failing at HEAD due to https://github.com/llvm/llvm-project/commit/72f10f7eb536da58cb79e13974895cd97d4e1a5f We can re-enable once https://github.com/llvm/llvm-project/commit/b9483d30a7d7a0650a0e83c75fcb9ab4932f475a is integrated. PiperOrigin-RevId: 601788984 26 January 2024, 17:36:07 UTC
3fc72d1 Fix `jax.lax.fori_loop(..., unroll=True)` with non-positive length 26 January 2024, 17:06:30 UTC
f34bcc3 Fixed a typo in Pallas GPU lowering `abs` is not available in `triton.compat.math`. PiperOrigin-RevId: 601709135 26 January 2024, 10:33:53 UTC
8901552 Update XLA dependency to use revision http://github.com/openxla/xla/commit/62156ca9ef4570b5dd2854eb9a93aaa9a4b23ca9. PiperOrigin-RevId: 601679131 26 January 2024, 07:37:10 UTC
70ea84d Merge pull request #19485 from ROCmSoftwarePlatform:rocm-enable_tridiagonal_solve PiperOrigin-RevId: 601613417 26 January 2024, 01:19:00 UTC
1264700 Merge pull request #19520 from jakevdp:fold-in-consume PiperOrigin-RevId: 601609582 26 January 2024, 01:01:28 UTC
b069c20 [key reuse] don't consume key in fold_in Why? We've found in practice that downstream projects use fold_in multiple times with the same key. This is safe so long as the folded-in value is different every time; in this sense fold_in() is similar to seed(), and for now we must trust the user to not repeat seeds. 25 January 2024, 23:35:51 UTC
45daced Merge pull request #19507 from jakevdp:wraps-implements PiperOrigin-RevId: 601505827 25 January 2024, 19:15:36 UTC
a6f2630 Update XLA dependency to use revision http://github.com/openxla/xla/commit/56977c4a88494b2618bbee8072901128bd38efdf. PiperOrigin-RevId: 601343380 25 January 2024, 06:42:37 UTC
fb80d2a [XLA][NFC] Make interface of module loaders consistent. LoadModuleFromData has (data, format, config, ...) signature while FromFile has (path, config, format, ...). Change the latter so `format` becomes the second argument in both cases. Since I'm touching this file: * Use `std::string_view` and `absl::Status` * Change `ovr_config` parameter to `const &` PiperOrigin-RevId: 601304308 25 January 2024, 03:16:43 UTC
5fb8b29 Remove `xla_gpu_cuda_data_dir` compile option from persistent cache key. See the new comment for justification. PiperOrigin-RevId: 601282907 25 January 2024, 01:26:59 UTC
ebfce19 Emit dense arrays for StableHLO ops migrating to dense arrays We are migrating some attrs on some StableHLO ops to use DenseI64ArrayAttr instead of DenseIntElementsAttr. Using DenseI64ArrayAttr enforces that the attr values are 1-dimensional and provides nicer APIs. (see https://github.com/openxla/stablehlo/issues/1578 for additional context) Unfortunately, we have to duplicate the `dense_int_array` function because we migrated the ops in batches. We can't use the existing `dense_int_array` function because it would produce arrays for ops that hadn't yet been migrated. This PR makes the final batch of changes, so no additional methods should be added going forward. We also have to introduce a new `dense_bool_array` function, with a similar version check. When the minimum supported jaxlib version uses a recent enough version of StableHLO (v6 or above), it will be possible to remove the version checks and remove the duplicated `dense_int_array_v6` function. PiperOrigin-RevId: 601271749 25 January 2024, 00:41:37 UTC
c7425ef Merge pull request #19509 from jakevdp:array-register PiperOrigin-RevId: 601264991 25 January 2024, 00:14:49 UTC
f15cad4 Lower a subset of math primitives directly to Triton IR Note that all primitives are now lowered to libdevice calls. Previously, some of them were lowered to the MLIR arith dialect, and some to libdevice calls, without any apparent reason for doing so. PiperOrigin-RevId: 601259707 24 January 2024, 23:55:09 UTC
78f27df Remove unnecessary Array.register 24 January 2024, 22:59:25 UTC
cfb6250 Add build instructions to build jaxlib with cuda plugin from source. PiperOrigin-RevId: 601231525 24 January 2024, 22:15:54 UTC
43a9faa Rename _wraps to implements 24 January 2024, 22:14:19 UTC
4646c64 [Pallas/TPU] Add support for input/output aliasing PiperOrigin-RevId: 601219571 24 January 2024, 21:37:19 UTC
8b81555 Merge pull request #19470 from jakevdp:device-arg PiperOrigin-RevId: 601205471 24 January 2024, 20:50:37 UTC
d55cd7c jax.numpy: support device argument for full, empty, zeros, ones 24 January 2024, 20:01:09 UTC
831c25f Merge pull request #19497 from ROCmSoftwarePlatform:rocm-log-pytest-html-report PiperOrigin-RevId: 601188997 24 January 2024, 19:51:28 UTC
b83f2b2 Merge pull request #19503 from jakevdp:np2-signatures PiperOrigin-RevId: 601177940 24 January 2024, 19:15:08 UTC
cedd67d Test: add weights to unsupported arguments 24 January 2024, 18:44:51 UTC
322f8b2 Merge pull request #19495 from 8bitmp3:patch-1 PiperOrigin-RevId: 601150181 24 January 2024, 17:48:29 UTC
ef7694f [ROCM]: Generating pytest html logs from unit-tests. 24 January 2024, 15:08:35 UTC
df3cc49 Update JAX docs copyright 24 January 2024, 14:19:45 UTC
a63197f Add an internal `_device_list` parameter to GSPMDSharding so that we can save on the initialization cost of PyDeviceList when creating `GSPMDSharding` from other shardings PiperOrigin-RevId: 601055733 24 January 2024, 10:29:22 UTC
a74b04a Merge pull request #19492 from gnecula:poly_tests PiperOrigin-RevId: 601050430 24 January 2024, 10:04:25 UTC
0bd511d [shape_poly] Add more tests for reasoning about inequalities. As I explore more powerful ways to reason about inequalities, I came up with more tests of inequalities that I wish we can handle. This PR adds the tests I have so far, even if they do not produce the correct result yet. I write the expected values for tests as _expect(best=v1, current=v2) to document that the current logic produces `v2` but the best value we can hope for is `v1`. This PR also adds more support for profiling tests. 24 January 2024, 08:57:49 UTC
eab0dd1 Update XLA dependency to use revision http://github.com/openxla/xla/commit/627d8c3f9d7c5c663d8d54ed107a084036786607. PiperOrigin-RevId: 601014027 24 January 2024, 06:52:12 UTC
6f96c96 Preserve single device NamedSharding/PositionalSharding on the output instead of always return SingleDeviceShardings. Fixes https://github.com/google/jax/issues/19459 PiperOrigin-RevId: 600999853 24 January 2024, 05:29:14 UTC
eb59716 Merge pull request #19145 from 8bitmp3:jax-docs-jax-primitives PiperOrigin-RevId: 600948053 24 January 2024, 00:56:09 UTC
a295b78 Upgrade JAX primitives doc 24 January 2024, 00:47:46 UTC
f67d482 Merge pull request #19489 from skye:version PiperOrigin-RevId: 600944466 24 January 2024, 00:41:11 UTC
368f594 Update version numbers for 0.4.24 release 24 January 2024, 00:36:10 UTC
88f5eac [xla:gpu] Make cu_threefry2x32 custom call compatible with command buffers PiperOrigin-RevId: 600937786 24 January 2024, 00:14:21 UTC
85fc4a9 Merge pull request #19484 from jakevdp:full-tests PiperOrigin-RevId: 600928191 23 January 2024, 23:39:28 UTC
037bc5e [Pallas] Add support for trivial vmap of scalar prefetch PiperOrigin-RevId: 600898742 23 January 2024, 21:57:12 UTC
fb4b491 [ROCM]: enable linalg_test.py::test_tridiagonal_solve. 23 January 2024, 21:20:18 UTC
8476b44 Test: move lax.full sharding tests to multi_device_test 23 January 2024, 20:10:00 UTC
7e80aac Merge pull request #19481 from gnecula:poly_changelog PiperOrigin-RevId: 600825615 23 January 2024, 18:04:54 UTC
af30668 [shape_poly] Forgot to update CHANGELOG for #19235. 23 January 2024, 16:32:59 UTC
c286372 Merge pull request #19473 from jakevdp:shard-map-run PiperOrigin-RevId: 600790090 23 January 2024, 16:02:09 UTC
8708d00 Merge pull request #19285 from gnecula:poly_constraints PiperOrigin-RevId: 600756934 23 January 2024, 13:29:04 UTC
5761e39 The Triton autotuner ignores configs that use too much shmem The autotuner runs a series of benchmarks to determine the best configuration for a Triton kernel. However, if it encounters a config that does not fit in shared memory it throws an error and stops. I this eventuality it should just continue. PiperOrigin-RevId: 600730507 23 January 2024, 11:08:57 UTC
24201ef [shape_poly] Add support for symbolic constraints on dimension variables Until now all the reasoning about symbolic dimensions was done with the implicit assumption that the dimension variables range over strictly positive integers. Here we allow the user to specify stronger constraints, so that they can be used in the reasoning about inequalities of symbolic dimensions. These explicit constraints are checked at compilation time, when the shapes are known. This adds significant power to the implementation of shape polymorphism, and in particular it adds an escape hatch for when in the past users saw inconclusive comparison exceptions. See more details in the README.md in this PR. 23 January 2024, 08:53:03 UTC
34d22fc Exclude modules with host callbacks from inter-host sharing. PiperOrigin-RevId: 600697796 23 January 2024, 08:38:13 UTC
f4b483c Update XLA dependency to use revision http://github.com/openxla/xla/commit/8398ac1000ba036f0dd5a15f0a49d9f13eb2b54f. PiperOrigin-RevId: 600669401 23 January 2024, 05:58:54 UTC
507239d Merge pull request #19474 from mattjj:shmap-disable-jit-fix PiperOrigin-RevId: 600657591 23 January 2024, 04:52:44 UTC
5275ef6 [shard-map] fix disable_jit-of-shmap 23 January 2024, 04:23:19 UTC
a3abd17 DOC: run shard_map docs as part of the doc build 23 January 2024, 02:07:01 UTC
70aec84 Merge pull request #19466 from jakevdp:full-like-sharding PiperOrigin-RevId: 600620827 23 January 2024, 01:48:13 UTC
f21022b Merge pull request #19472 from mattjj:shmap-tutorial-typos PiperOrigin-RevId: 600619148 23 January 2024, 01:40:15 UTC
1df92c9 Merge pull request #19438 from sharadmv:add-init PiperOrigin-RevId: 600618332 23 January 2024, 01:32:21 UTC
a61f780 [shard-map] fix typos in docs 23 January 2024, 01:29:10 UTC
76554bc [Pallas/TPU] Add vmem limit test PiperOrigin-RevId: 600571425 22 January 2024, 22:34:14 UTC
dfda948 [XLA:Python] Fail with an AttributeError if __cuda_array_interface__ is called on a sharded array. Fixes https://github.com/google/jax/issues/19134 PiperOrigin-RevId: 600570354 22 January 2024, 22:25:47 UTC
46f796b Dedupe shardings before passing them to _get_and_check_device_assignment In practice, the number of different shardings is usually much smaller then the number of inputs/output. PiperOrigin-RevId: 600558309 22 January 2024, 21:45:20 UTC
8226ff3 Merge pull request #19455 from ROCmSoftwarePlatform:rocm-jaxlib-docker-rocm6.0-updates PiperOrigin-RevId: 600553582 22 January 2024, 21:28:50 UTC
71ee3ab Disable NumPy interoperability DLPack test for bools for NumPy versions older than 1.25. DLPack bool support was added in NumPy 1.25. Fixes a test failure in CI. PiperOrigin-RevId: 600531477 22 January 2024, 20:13:49 UTC
25e4acf Disabled lax_scipy_special_functions_test under ASAN on GPU. This test is slow and times out in CI. PiperOrigin-RevId: 600527658 22 January 2024, 20:02:34 UTC
e76b87e lax.full_like: add sharding argument 22 January 2024, 19:55:25 UTC
a415b56 Merge pull request #19464 from jakevdp:actions-cache PiperOrigin-RevId: 600512505 22 January 2024, 19:21:30 UTC
9b9aa1e Finalize a number of deprecations from JAX 0.4.19 PiperOrigin-RevId: 600509530 22 January 2024, 19:13:25 UTC
416206b Merge pull request #19460 from google:dependabot/github_actions/actions/upload-artifact-4.2.0 PiperOrigin-RevId: 600508301 22 January 2024, 19:05:22 UTC
ca63581 Merge pull request #19463 from jakevdp:numpy-upstream PiperOrigin-RevId: 600508042 22 January 2024, 18:56:53 UTC
3392b64 Bump actions/cache from 3.3.3 to 4.0.0 22 January 2024, 18:54:32 UTC
34901b2 TST: add unimplemented parameters in numpy signatures test 22 January 2024, 18:43:52 UTC
e23ee15 Merge pull request #19445 from jakevdp:full-device PiperOrigin-RevId: 600502184 22 January 2024, 18:39:27 UTC
c12cf55 Merge pull request #19462 from superbobry:legacy-directive PiperOrigin-RevId: 600496307 22 January 2024, 18:20:54 UTC
4196bf1 DOC Add a noop implementation of the legacy directive to fix the build 22 January 2024, 17:59:54 UTC
577230e Bump actions/upload-artifact from 4.1.0 to 4.2.0 Bumps [actions/upload-artifact](https://github.com/actions/upload-artifact) from 4.1.0 to 4.2.0. - [Release notes](https://github.com/actions/upload-artifact/releases) - [Commits](https://github.com/actions/upload-artifact/compare/1eb3cb2b3e0f29609092a73eb033bb759a334595...694cdabd8bdb0f10b2cea11669e1bf5453eed0a6) --- updated-dependencies: - dependency-name: actions/upload-artifact dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> 22 January 2024, 17:51:38 UTC
f41a32c lax.full: add sharding argument 22 January 2024, 17:27:47 UTC
f997609 [ROCm]: Updates hip headers path for ROCm 6.0 22 January 2024, 16:08:37 UTC
b7a7f0b [ROCm]: Dockerfile updates 22 January 2024, 16:08:37 UTC
b512b57 Update XLA dependency to use revision http://github.com/openxla/xla/commit/191cf785dc92a8e1df53988761f8dad7867f78dc. PiperOrigin-RevId: 600337284 22 January 2024, 05:07:22 UTC
back to top