ROCm/jax

Pull Requests Commits

Temporarily roll back changes for new LLVM version

slackito committed 1 year ago

51cd3c3b

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

junwhanahn committed 1 year ago

35b25c91

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

Google-ML-Automation committed 1 year ago

80fc4fab

Merge pull request #28507 from findmyway:patch-1

Google-ML-Automation committed 1 year ago

3299aed3

Merge pull request #28533 from chaserileyroberts:chaser/multiprocess_test_fix

Google-ML-Automation committed 1 year ago

40f1ddfd

Integrate LLVM at llvm/llvm-project@2d287f51eff2

alinas committed 1 year ago

439f0f14

Generalized test, removed specific call to config

chaserileyroberts committed 1 year ago

8c113bd1

Avoid an unlucky seed for for some random categorical tests.

jburnim committed 1 year ago

cf9b265c

Merge pull request #28485 from andportnoy:patch-5

Google-ML-Automation committed 1 year ago

64e6f93f

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

superbobry committed 1 year ago

579a1573

Enable command buffer support for buffer callbacks.

dfm committed 1 year ago

958ea15a

Fix a x64 error in fused_attention_stablehlo

apaszke committed 1 year ago

9c65f62d

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

b3c49b0e

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

apaszke committed 1 year ago

db376303

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

dfm committed 1 year ago

72e1a7d2

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

apaszke committed 1 year ago

70977f08

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

apaszke committed 1 year ago

66fc5e88

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

apaszke committed 1 year ago

cd513f25

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

emilyfertig committed 1 year 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 1 year ago

3193ed39

PR #28157: Warn on excessive captured constants

ZacCranko committed 1 year ago

d9aee152

Update sharded-computation.ipynb

findmyway committed 1 year ago

Verified 47b13c77

[docs] Remove tensorflow from tensorboard installation instructions

rdyro committed 1 year 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 1 year ago

6d1b5271

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

Jake VanderPlas committed 1 year ago

cf976283

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

vfdev-5 committed 1 year ago

2c52ae35

Add workflow for testing JAX against NumPy nightly builds.

nitins17 committed 1 year 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 1 year 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 1 year ago

a2d67a6a

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

Google-ML-Automation committed 1 year ago

52fcf84a

Older