ROCm/jax

Pull Requests Commits

Skip CSR matmat/matvec float tests on ROCm <6.4 (NaN issue with beta==0).

psanal35 committed 1 year ago

0af3e820

Rename `SdyArraySharding -> SdyArray` and `SdyDimSharding -> SdyDim` since these are not `Sharding` from a JAX POV.

yashk2810 committed 1 year ago

3142bc3b

Replace gsutil command with gcloud storage commands.

nitins17 committed 1 year ago

abc79f3d

[Pallas/Mosaic GPU] Allow `_is_known_divisible` to always return `True` when the divisor is `1`.

bchetioui committed 1 year ago

85a68fcd

[pallas:mosaic_gpu] Do not do unnecessary `commit_smem_to_gmem_group` in `emit_pipeline`

superbobry committed 1 year ago

af66ca9c

[Pallas/Mosaic GPU] Optimize the construction of output buffers for `plgpu.kernel`.

bchetioui committed 1 year ago

7f5b6e7d

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

50970c6b

Simplify if_building_jaxlib macro.

hawkinsp committed 1 year ago

67e1d5c2

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

3d64dd78

allow accelerated deprecations from jax.tree_util

Jake VanderPlas committed 1 year ago

dfb5a93f

Prepare to make DmaCopyChunk movable.

pschuh committed 1 year ago

9b2043e4

Global find/replace in jax for s/divisble/divisible.

jkr26 committed 1 year ago

e2c4a7f5

Deprecate the no-op custom_jvp_call_jaxpr_p import stub.

dfm committed 1 year ago

c7a4e34b

Pass RULES_PYTHON_REPO_DEBUG value to the docker container env

kanglant committed 1 year ago

49e3432b

Merge pull request #28722 from jakevdp:jaxlib-api-deprecation

Google-ML-Automation committed 1 year ago

e6b47535

[Mosaic TPU] Fold reshape (..., M, N, 128) -> (..., M, N * 128) with any tiling and dtype to load.

bythew3i committed 1 year ago

7887298d

Merge pull request #28661 from jakevdp:jax-array-jep

Google-ML-Automation committed 1 year ago

0c10daed

Avoid duplication of Bazel dependencies between //jax/_src/lib and //jaxlib.

hawkinsp committed 1 year ago

b2549761

[doc] mention jaxlib in the API compatibility doc

jakevdp committed 1 year ago

6971e125

Fix a stale type annotation.

hawkinsp committed 1 year ago

d89f2a56

Merge pull request #28355 from jakevdp:jax-array-abstract

Google-ML-Automation committed 1 year ago

e55dadd2

Testing: change jax.extend.ffi to jax.ffi

Jake VanderPlas committed 1 year ago

a98f533c

Add direct `pypi` dependencies to the JAX test targets.

hawkinsp committed 1 year ago

725d0f64

Use DmaCopyChunk::Make because directly assigning the struct

pschuh committed 1 year ago

0e4f213e

[JAX] Make fully replicated sharding to avoid materializing the same host buffers

hyeontaek committed 1 year ago

ef1b3e92

Merge pull request #28715 from dfm:custom-vjp-pp

Google-ML-Automation committed 1 year ago

d6608988

[Pallas] Allow f8 casting tests on TPUv5-.

WindQAQ committed 1 year ago

78f89b8b

Add a pretty printing rule for custom_vjp.

dfm committed 1 year ago

71692fcb

Fix xla_bridge_test: Do "assert subset" the right way around.

matthiaskramm committed 1 year ago

60a37718

Merge pull request #28711 from dfm:custom-vjp-symb-zeros

Google-ML-Automation committed 1 year ago

8060ca2e

Older