ROCm/jax

Pull Requests Commits

Merge pull request #30446 from jakevdp:source-info-util-build

Google-ML-Automation committed 182 days ago

9b47a902

Cache `canonicalize_sharding` and remove the canonicalization in `_scatter_update` since the out_sharding passed to scatter_update should already be

yashk2810 committed 182 days ago

768a64eb

avoid top-level imports in source_info_util

jakevdp committed 182 days ago

338b4ebc

Skip Pallas and Mosaic GPU tests that don't fit on RTX 6000 PRO

andportnoy committed 182 days ago

240c2d0a

Test B200 against CUDA 12.8 only

nitins17 committed 182 days ago

fd8a79f8

Add Hermetic C++ Toolchains for Linux x86_64 builds.

Google-ML-Automation committed 182 days ago

bf25960c

Merge pull request #30257 from andportnoy:aportnoy/mosaic-gpu-query-shared-memory

Google-ML-Automation committed 182 days ago

b2bd5e77

[jax:ffi] Expose `build_ffi_lowering_function` to give user access to the lowering function that returns an Op, not just Op results.

danielsuo committed 182 days ago

e381a868

[Pallas TPU] Remove `use_sreg_for_state` flag

Google-ML-Automation committed 182 days ago

4785804f

Remove out_sharding from scatter_add because we now have a scatter sharding rule

yashk2810 committed 182 days ago

1c0f24b2

Merge pull request #30440 from jax-ml:direct-linearize-context-bug

Google-ML-Automation committed 182 days ago

d5898929

Use the correct tracing context when creating zeros in linearize_from_jvp

dougalm committed 182 days ago

88eeb4fb

Implement explicit mode sharding rule for scatter

oliverdutton-iso committed 182 days ago

23563247

#sdy Add unit tests making sure we fallback to GSPMD if the original model was exported with GSPMD but loaded with Shardy.

bartchr808 committed 182 days ago

cb10a759

[jax] Export ir_attribute in jax.interpreters.mlir

James Molloy committed 182 days ago

96c7d61e

Fix typo in doc.

allanrenucci committed 182 days ago

41afeac1

[Mosaic] NFC: Just use ExtFOp and TruncFOp for bf16.

WindQAQ committed 182 days ago

9230c676

[Mosaic GPU] Simplify type constraints in MGPU dialect.

allanrenucci committed 182 days ago

8bb6b00f

[Mosaic GPU] Add support for warp shuffles with elements wider than 32-bit

apaszke committed 182 days ago

844a644b

Fix a bug where doing a CPU jax.Array -> TPU jax.Array transfer was not setting the layout correctly on the output TPU array.

yashk2810 committed 183 days ago

b57f0773

Update XLA dependency to use revision

Google-ML-Automation committed 183 days ago

7a6f1dfe

[Mosaic] Add tpu.wait_indirect_dma assembly format

naummo committed 183 days ago

7f1934c3

Merge pull request #30132 from carlosgmartin:numpy_size_multiple_axes

Google-ML-Automation committed 183 days ago

f8e8237d

Avoid top-level imports in jax/_src/internal_test_util/*

Jake VanderPlas committed 183 days ago

a350bde9

Add `portpicker` to the list of common requirements and update `requirements_lock_3_13_ft.txt` and `requirements_lock_3_14_ft.txt` files.

Google-ML-Automation committed 183 days ago

e76eb183

Avoid top-level imports in jax/_src/extend

Jake VanderPlas committed 183 days ago

da7d0ff5

Remove JAX_IFRT_VERSION_NUMBER guards after the release of 0.7.0

yashk2810 committed 183 days ago

fb1a8fcb

#sdy Delete JAX test configs enabling Shardy.

yashk2810 committed 183 days ago

ff595f61

Refactoring HTTIJ and surrounding pages

melissawm committed 183 days ago

f99ba56b

Add LRU cache on `get_process_index_and_count`.

Hilly12 committed 183 days ago

75229a69

Newer Older