ROCm/jax

Pull Requests Commits

allow accelerated deprecations from jax.tree_util

Jake VanderPlas committed 228 days ago

dfb5a93f

Prepare to make DmaCopyChunk movable.

pschuh committed 229 days ago

9b2043e4

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

jkr26 committed 229 days ago

e2c4a7f5

Deprecate the no-op custom_jvp_call_jaxpr_p import stub.

dfm committed 229 days ago

c7a4e34b

Pass RULES_PYTHON_REPO_DEBUG value to the docker container env

kanglant committed 229 days ago

49e3432b

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

Google-ML-Automation committed 229 days ago

e6b47535

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

bythew3i committed 229 days ago

7887298d

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

Google-ML-Automation committed 229 days ago

0c10daed

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

hawkinsp committed 229 days ago

b2549761

[doc] mention jaxlib in the API compatibility doc

jakevdp committed 229 days ago

6971e125

Fix a stale type annotation.

hawkinsp committed 229 days ago

d89f2a56

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

Google-ML-Automation committed 229 days ago

e55dadd2

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

Jake VanderPlas committed 229 days ago

a98f533c

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

hawkinsp committed 229 days ago

725d0f64

Use DmaCopyChunk::Make because directly assigning the struct

pschuh committed 229 days ago

0e4f213e

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

hyeontaek committed 229 days ago

ef1b3e92

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

Google-ML-Automation committed 229 days ago

d6608988

[Pallas] Allow f8 casting tests on TPUv5-.

WindQAQ committed 229 days ago

78f89b8b

Add a pretty printing rule for custom_vjp.

dfm committed 229 days ago

71692fcb

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

matthiaskramm committed 229 days ago

60a37718

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

Google-ML-Automation committed 229 days ago

8060ca2e

Don't instantiate zeros passed to custom_lin_p.

dfm committed 229 days ago

df66c2fd

Merge pull request #28690 from jakevdp:revert-28650

Google-ML-Automation committed 229 days ago

f62b7900

[Mosaic] Make `tpu.relayout` an explicit operation, merge in existing behavior, stop calling relayout() in apply.

Google-ML-Automation committed 229 days ago

123022ca

Merge pull request #28589 from dfm:consolidate-custom-vjp-primitives

Google-ML-Automation committed 229 days ago

1ad9eae3

[Mosaic GPU] Add layout inference and lowering for `scf.WhileOp` and enable tests.

dimitar-asenov committed 229 days ago

66476a0b

Update XLA dependency to use revision

Google-ML-Automation committed 229 days ago

f73726e9

[pallas:mosaic_gpu] Nuked a debug test added by accident

superbobry committed 229 days ago

0d6ad8ae

[pallas:mosaic_gpu] Removed debug prints from `emit_pipeline_warp_specialized`

superbobry committed 229 days ago

3667353d

Update tridiagonal solve kernels on GPU to properly use the FFI.

dfm committed 229 days ago

91de2e39

Older