ROCm/jax

Pull Requests Commits

[Pallas/interpreter] Support interpreting non-trivial GPU grids.

Google-ML-Automation committed 37 days ago

d799c9cc

Expand and update XProf docs

mtsokol committed 37 days ago

335bd052

[Pallas:MGPU] Increase shard_count for the interpret test

apaszke committed 37 days ago

3ebe54ed

Make `sub_byte_element_size_in_bits` public

Google-ML-Automation committed 37 days ago

2acc272a

Merge pull request #31720 from ROCm:ann-test-fix

Google-ML-Automation committed 37 days ago

81d95401

[Pallas:SC] Skip a debug_check test failing on TPU7x

apaszke committed 37 days ago

e5aad81b

Update XLA dependency to use revision http://github.com/openxla/xla/commit/59a18ff2183fd80a0b0af123eff42cd85ba72bb8

Google-ML-Automation committed 37 days ago

b1272723

[pmap] Speed up the new `jax.pmap` implementation.

danielsuo committed 37 days ago

cf763131

[pmap] Optimize host-local <> global conversions.

danielsuo committed 37 days ago

a2912b07

Fix: Use ASAN feature instead of config setting

yuriivcs committed 37 days ago

678acf28

[pmap] Push more into _cached_shard_map.

danielsuo committed 37 days ago

ee911642

[Mosaic] Expose `tpu::isGather`

naummo committed 38 days ago

2dbb22dd

Add env variables JAX_TRITON_PTXAS_DEVICE_DEBUG JAX_TRITON_PTXAS_EXTRA_FLAGS.

Google-ML-Automation committed 38 days ago

1f12c0a2

[mosaic:gpu] Use bytecode serialization.

hawkinsp committed 38 days ago

abae5e4e

[pmap] Remove any wrappings that have auxiliary values.

danielsuo committed 38 days ago

49561f67

[pmap] Cache computing effective mesh devices.

danielsuo committed 38 days ago

1cf66905

Only consider Tracers with ShapedArray avals as instances of Array.

hawkinsp committed 38 days ago

1b6e0981

[export] Remove serialization of old 16-bit nr_devices

gnecula committed 38 days ago

9f327f69

Fix compute_on2 with out_memory_spaces == Host

yashk2810 committed 38 days ago

191504b6

Include source info annotations in state primitive custom pretty printer rules.

hawkinsp committed 38 days ago

6dc73528

Merge pull request #34644 from ROCm:rocm-skip-prim-tri-solve-test

Google-ML-Automation committed 38 days ago

3a38a4bf

[pmap] Inline _prepare_pmap and clean up unused structs.

danielsuo committed 38 days ago

692d820f

Update XLA dependency to use revision http://github.com/openxla/xla/commit/e9b87bec13cef64ca946f6e79a97c6c9ddf16203

Google-ML-Automation committed 38 days ago

762f0661

[pmap] Remove check_arg on every pmap invocation.

danielsuo committed 38 days ago

08fd5ea3

Switch XLA archive download URL from GitHub API to direct archive.

nitins17 committed 38 days ago

b839e6c3

Make `halt-for-connection` passable from workflow call.

ybaturina committed 38 days ago

44cc4f85

[ROCm] Added skip for "test_prim_tridiagonal_solve" tests on ROCm.

tsrw2048 committed 38 days ago

5eb5a215

[pmap] Add fast-paths for inferring and mapping in_axes.

danielsuo committed 38 days ago

d1fff7f5

[Mosaic GPU][NFC] Pass `CompiledKernel` by pointer to `CachedInit`.

allanrenucci committed 38 days ago

745cbe2f

[Pallas/Mosaic GPU] Use `self.capture_stdout` instead of `jtu.capture_stdout` in `mosaic_gpu_test`.

bchetioui committed 38 days ago

7241ee09

Newer Older