Jake VanderPlas
c7b0d681bd
Remove deprecated jax.experimental.array_api
2025-01-06 15:19:02 -08:00
jax authors
52cc5c7f05
Merge pull request #25214 from jakevdp:einsum-optimize
...
PiperOrigin-RevId: 712603103
2025-01-06 11:37:54 -08:00
Jevin Jiang
9f842909ce
[Mosaic TPU] Validate inserted layout in relayout-insertion pass.
...
PiperOrigin-RevId: 712595778
2025-01-06 11:15:47 -08:00
jax authors
634b45bf00
Merge pull request #25699 from yliu120:fix_iota
...
PiperOrigin-RevId: 712594991
2025-01-06 11:13:28 -08:00
Jake VanderPlas
2f7204fff6
jnp.einsum: default to optimize='auto'
2025-01-06 11:02:31 -08:00
John QiangZhang
c39e38fe5a
bazel: export serialization.fbs for downstream usage
...
PiperOrigin-RevId: 712587802
2025-01-06 10:57:35 -08:00
jax authors
74be8bd99f
Merge pull request #25675 from jakevdp:dep-lpmn
...
PiperOrigin-RevId: 712579230
2025-01-06 10:37:10 -08:00
jax authors
18b193cbbd
Update XLA dependency to use revision
...
1a6361a734
2025-01-06 09:41:02 -08:00
Jake VanderPlas
245a13a329
Deprecate scipy.special.lpmn & lpmn_values
2025-01-06 09:31:15 -08:00
Mark Sandler
6c87bf389f
Fixes tril/triu comments (they were flipped)
...
PiperOrigin-RevId: 712544847
2025-01-06 08:55:11 -08:00
Yunlong Liu
3ff000ee3e
fix the degenerated case
2025-01-06 16:08:07 +00: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
