Mark Sandler
6c87bf389f
Fixes tril/triu comments (they were flipped)
...
PiperOrigin-RevId: 712544847
2025-01-06 08:55:11 -08:00
George Necula
e87a2a5929
[shape_poly] Remove old non_negative support.
...
This was deprecated in January 2024, replaced by
`core_max_dim(..., 0)`.
PiperOrigin-RevId: 712523579
2025-01-06 07:36:11 -08:00
Vladimir Belitskiy
f2e210b315
Disable avxvnniint8 when building with Clang version < 19, or GCC < 13.
...
PiperOrigin-RevId: 712516025
2025-01-06 07:06:09 -08:00
Dan Foreman-Mackey
512d5450ae
Temporarily allow deprecation warnings for scipy.special.lpmn and scipy.special.sph_harm.
...
These functions are deprecated in scipy 1.15.0. I'll fix this properly soon, but let's start by getting CI working again!
PiperOrigin-RevId: 712512363
2025-01-06 06:50:51 -08:00
jax authors
4de7794faf
Merge pull request #25715 from ROCm:ci_build_code_fixes-upstream
...
PiperOrigin-RevId: 712463302
2025-01-06 03:05:14 -08:00
jax authors
d0a92c5c7d
Update XLA dependency to use revision
...
ac6e71fe0c
2025-01-05 08:45:44 -08:00
jax authors
54fd738ecb
Add SMEM as a supported Pallas output memory space.
...
PiperOrigin-RevId: 712144883
2025-01-04 19:33:18 -08:00
jax authors
9af2970042
Update XLA dependency to use revision
...
c12c114858
2025-01-04 08:55:59 -08:00
jax authors
e4278f7866
Update XLA dependency to use revision
...
a84e3b7f8f
2025-01-03 09:07:45 -08:00
jax authors
0f4677bcf6
Merge pull request #25713 from jakevdp:debug-printoptions
...
PiperOrigin-RevId: 711671926
2025-01-03 01:33:31 -08:00
Jake VanderPlas
330606320a
jax.debug.print: respect local np.printoptions
2025-01-02 16:10:54 -08:00
Tzu-Wei Sung
57b21541a2
[Mosaic] NFC: Pull out vreg related functions to util.
...
These functions are related to vreg manipulation and are used in different rules.
PiperOrigin-RevId: 711484002
2025-01-02 11:50:19 -08:00
Zac Mustin
df36c29803
Compute cost-analysis on only one HLO module.
...
There was historically a goal to support multiple HLOs in an executable, but this work was never finished and is no longer planned so we don't need this support.
This will soon enable us to return only a dict, instead of a list of dicts with only one item.
PiperOrigin-RevId: 711477481
2025-01-02 11:24:52 -08:00
jax authors
800f903f9b
Merge pull request #25686 from Mikcl:docs/working-with-pytrees-formatting
...
PiperOrigin-RevId: 711448394
2025-01-02 09:40:17 -08:00
jax authors
726950b885
Update XLA dependency to use revision
...
06078480db
2025-01-02 08:48:53 -08:00
jax authors
68483b8ed6
Merge pull request #25710 from apaszke:mgpu_dialect_fix
...
PiperOrigin-RevId: 711430610
2025-01-02 08:23:28 -08:00
Adam Paszke
64433435ff
Fix OSS build for the Mosaic GPU dialect
2025-01-02 15:55:03 +00:00
Tomás Longeri
ac817b48ca
[Mosaic:TPU][NFC] Clean up unused variable
...
PiperOrigin-RevId: 711412888
2025-01-02 06:57:38 -08:00
jax authors
82001ed5b3
Merge pull request #25706 from pearu:pearu/log10-large
...
PiperOrigin-RevId: 711411578
2025-01-02 06:50:54 -08:00
jax authors
04a0fbecc3
Merge pull request #25661 from rdyro:tb-nightly-instructions
...
PiperOrigin-RevId: 711407514
2025-01-02 06:27:24 -08:00
Adam Paszke
7c984c600b
Don't use x32 mode for pallas_test
...
There's no need to, and it caused our GPU tests for this target to only
run nightly.
PiperOrigin-RevId: 711406571
2025-01-02 06:23:32 -08:00
Adam Paszke
dbe9ccd6dc
Reverts 83e60a9697ec20023f4e11169edf64e910b93031
...
PiperOrigin-RevId: 711403091
2025-01-02 06:04:14 -08:00
Robert Dyro
213e1782ac
tbp nightly instructions
2025-01-02 09:49:31 +01:00
Ruturaj4
20b75ab82f
Update package indentation fix
2025-01-01 18:50:47 -06:00
jax authors
4a6cfebcea
Update XLA dependency to use revision
...
045356d8c8
2025-01-01 08:23:24 -08:00
Pearu Peterson
50670bd907
Fix log10 and log2 for large inputs.
2025-01-01 12:45:39 +02:00
jax authors
e37ea587c9
Update XLA dependency to use revision
...
cb0c24822f
2024-12-31 08:49:31 -08:00
jax authors
494c15733c
Merge pull request #25692 from yliu120:rng_key_sharding
...
PiperOrigin-RevId: 710742320
2024-12-30 10:10:40 -08:00
jax authors
25fff524f0
Update XLA dependency to use revision
...
53e42f4ecd
2024-12-30 09:02:36 -08:00
Yunlong Liu
97b1faacdd
Fixes the random key sharding in shard_map.
2024-12-29 18:43:21 +00:00
jax authors
879fa12d90
Update XLA dependency to use revision
...
54ce5aecdc
2024-12-29 09:53:34 -08:00
jax authors
8eeedd1802
Update XLA dependency to use revision
...
50420127c5
2024-12-28 10:26:46 -08:00
Sergei Lebedev
76ccb199fd
[pallas:mosaic_gpu] Added some runtime type checking to copy_* and barrier_* primitives
...
PiperOrigin-RevId: 710302436
2024-12-28 09:02:43 -08:00
jax authors
7ab61b7f26
Update XLA dependency to use revision
...
15efc84a0c
2024-12-27 10:57:06 -08:00
jax authors
6dbda900d1
Update XLA dependency to use revision
...
e52f275aba
2024-12-26 11:03:35 -08:00
mikcl
008c25a369
Fix formatting in the docs for transposing pytrees
2024-12-26 01:18:10 +00:00
jax authors
42a0d55503
Update XLA dependency to use revision
...
e268cb7c7f
2024-12-25 10:21:10 -08:00
jax authors
b6aead6f3a
[AutoPGLE] Explicitly disable command buffers when profiler is used.
...
PiperOrigin-RevId: 709475833
2024-12-24 21:31:05 -08:00
jax authors
64511a19fd
Update XLA dependency to use revision
...
16e5b0cf64
2024-12-24 10:05:12 -08:00
Sergei Lebedev
44333e1cfb
[pallas:mosaic_gpu] Addressed a todo in broadcasted_iota lowering
...
PiperOrigin-RevId: 709310152
2024-12-24 04:32:29 -08:00
jax authors
4eff1316db
Merge pull request #25672 from jakevdp:finalize-dep
...
PiperOrigin-RevId: 709284584
2024-12-24 02:08:51 -08:00
jax authors
fa9c7edf73
Merge pull request #25674 from jax-ml:dependabot/github_actions/actions/upload-artifact-4.5.0
...
PiperOrigin-RevId: 709190327
2024-12-23 17:39:01 -08:00
Tomás Longeri
4452960947
[Mosaic:TPU] In infer ext rule, avoid assigning offsets outside of dst first tile
...
Note that offsets outside of first tile are still disabled (for both infer and apply), and once we support it we will want to assign offsets differently, this is mostly to avoid assigning invalid layouts (that may not just be outside the first tile, but outside the vreg slice)
PiperOrigin-RevId: 709168368
2024-12-23 15:49:39 -08:00
jax authors
b8091a437a
Switch mlir bindings from pybind11 to nanobind
...
PiperOrigin-RevId: 709161113
2024-12-23 15:10:11 -08:00
Tomás Longeri
3c79b98cd9
[Mosaic:TPU] Vreg-slice-aligned offset changes with scratch retiling
...
PiperOrigin-RevId: 709133729
2024-12-23 13:05:14 -08:00
jax authors
c57b49c606
Merge pull request #25669 from jakevdp:undep
...
PiperOrigin-RevId: 709104784
2024-12-23 11:06:24 -08:00
Jake VanderPlas
40fe4b8797
Finalize deprecation of some symbols from jax.lib.xla_client
2024-12-23 10:14:16 -08:00
dependabot[bot]
7da753eb0b
Bump actions/upload-artifact from 4.4.3 to 4.5.0
...
Bumps [actions/upload-artifact](https://github.com/actions/upload-artifact ) from 4.4.3 to 4.5.0.
- [Release notes](https://github.com/actions/upload-artifact/releases )
- [Commits](b4b15b8c7c...6f51ac03b9
2024-12-23 18:01:03 +00:00
jax authors
108762379d
Merge pull request #25673 from jakevdp:changelog
...
PiperOrigin-RevId: 709086770
2024-12-23 09:55:42 -08:00
Jake VanderPlas
c206ae7fe8
changelog: link to api compatibility & python version docs
2024-12-23 09:39:45 -08:00
