ROCm/jax

Pull Requests Commits

Temporarily roll back changes for new LLVM version

slackito committed 261 days ago

51cd3c3b

Use `LoadedExecutableRef` instead of `std::unique_ptr<LoadedExecutable>`

junwhanahn committed 261 days ago

35b25c91

Merge pull request #28387 from vfdev-5:fix-deadlock-makeshardfn-28385

Google-ML-Automation committed 261 days ago

80fc4fab

Merge pull request #28507 from findmyway:patch-1

Google-ML-Automation committed 261 days ago

3299aed3

Merge pull request #28533 from chaserileyroberts:chaser/multiprocess_test_fix

Google-ML-Automation committed 261 days ago

40f1ddfd

Integrate LLVM at llvm/llvm-project@2d287f51eff2

alinas committed 261 days ago

439f0f14

Generalized test, removed specific call to config

chaserileyroberts committed 261 days ago

8c113bd1

Avoid an unlucky seed for for some random categorical tests.

jburnim committed 261 days ago

cf9b265c

Merge pull request #28485 from andportnoy:patch-5

Google-ML-Automation committed 261 days ago

64e6f93f

[pallas:mosaic] Handle more types in `ir_constant`

superbobry committed 261 days ago

579a1573

Enable command buffer support for buffer callbacks.

dfm committed 261 days ago

958ea15a

Fix a x64 error in fused_attention_stablehlo

apaszke committed 261 days ago

9c65f62d

Update XLA dependency to use revision

Google-ML-Automation committed 261 days ago

b3c49b0e

[Mosaic-GPU] [3/3] Add support for communication primitives in MGPU lowering

apaszke committed 261 days ago

db376303

Enable batch sharding tests for Cholesky and triangular solve on GPU.

dfm committed 261 days ago

72e1a7d2

[Mosaic GPU] Add tsl::TraceMe to annotate compilation cache misses

apaszke committed 262 days ago

70977f08

[Mosaic GPU] Add support for arbitrary reductions of tiled layouts

apaszke committed 262 days ago

66fc5e88

[Pallas:MGPU] Perform a warpgroup barrier before and after every SMEM write

apaszke committed 262 days ago

cd513f25

Remove obsolete DCHECK on num_computations_, now that executables may have no local devices.

emilyfertig committed 262 days ago

766e68c4

Move attrs.py code to `_src/attrs.py` so that we can remove experimental imports in pjit.py which are in the middle of tracing code and can affect performance.

yashk2810 committed 262 days ago

3193ed39

PR #28157: Warn on excessive captured constants

ZacCranko committed 262 days ago

d9aee152

Update sharded-computation.ipynb

findmyway committed 262 days ago

Verified 47b13c77

[docs] Remove tensorflow from tensorboard installation instructions

rdyro committed 262 days ago

8f17a552

Support `np.asarray` for JAX arrays with fully replicated shardings when some hosts have no local shards. Only hosts with local shards can fetch np arrays.

emilyfertig committed 262 days ago

6d1b5271

Replace C++ safe_zip with Python implementation using 3.10 strict keyword

Jake VanderPlas committed 262 days ago

cf976283

Fixed deadlock in MakeShardFn on static var assignment under free-threading

vfdev-5 committed 262 days ago

2c52ae35

Add workflow for testing JAX against NumPy nightly builds.

nitins17 committed 262 days ago

13ca7002

[JAX] In McJAX, support JIT compilation/execution on only the devices attached to a single host, with other hosts not participating.

emilyfertig committed 262 days ago

268c6842

create a mosaic_gpu_support cc_library and include it as part of the //third_party/py/jax/jaxlib/cuda:cuda_gpu_kernels

Google-ML-Automation committed 262 days ago

a2d67a6a

Merge pull request #28513 from olupton:skip-cublas-12.9

Google-ML-Automation committed 262 days ago

52fcf84a

Older