ROCm/jax

Pull Requests Commits

Inject rocm plugin deps

alekstheod committed 228 days ago

9e0cb6b0

Migrate rocm-jax repo changes back to upstream jax

alekstheod committed 232 days ago

d9bc1c0c

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

Google-ML-Automation committed 233 days ago

aee33099

Add memory_space to typematch check

yashk2810 committed 233 days ago

540df654

Merge pull request #35095 from carlosgmartin:lax_reduce_doc_init_value_scalar

Google-ML-Automation committed 233 days ago

4a374b41

Merge pull request #35153 from davisyoshida:hipspec

Google-ML-Automation committed 233 days ago

7ecbc030

Export HipSpec from hijax

davisyoshida committed 233 days ago

fc641811

Merge pull request #35139 from jakevdp:fix-lint

Google-ML-Automation committed 233 days ago

ab7aecd7

Merge pull request #35131 from jakevdp:scan-length-check

Google-ML-Automation committed 233 days ago

80ebaef8

Restore scan length check

jakevdp committed 233 days ago

1eed1171

Reverts 4a14e0c2cce9ff41cbca271b429d36df8215e608

nvgrw committed 233 days ago

5abf880a

Fix jupytext lint at head

jakevdp committed 233 days ago

9b317766

Merge pull request #35137 from mattjj:issue34758

Google-ML-Automation committed 233 days ago

81640be3

Merge pull request #35136 from mattjj:hi-vmap

Google-ML-Automation committed 233 days ago

1c1ad14e

[hijax] test vmap + hijax types + hijax primitive

mattjj committed 233 days ago

81ebe31f

[hijax] add cur_qdd to BatchTrace

mattjj committed 233 days ago

b456d935

Merge pull request #35083 from mattjj:hi-shmap

Google-ML-Automation committed 234 days ago

0ea5ee02

[hijax] basic shmap-of-hitypes by adding {,un}shard_aval and HipSpec

mattjj committed 234 days ago

cf879e1b

Integrate LLVM at llvm/llvm-project@c6bb1afbc508

Google-ML-Automation committed 234 days ago

6a0d96d3

[pallas:sc] The default memory space in `aval_to_ir_type` now depends on the kernel type

superbobry committed 234 days ago

70692c6f

Disable an einshape test on TPU 7x that is failing

hawkinsp committed 234 days ago

696a10ae

Merge pull request #35087 from jakevdp:more-pyrefly-fixes

Google-ML-Automation committed 234 days ago

a10ee0ae

[Mosaic GPU][NFC] Remove duplicated python binding.

dimitar-asenov committed 234 days ago

e0b75368

[pyrefly] more typing fixes

jakevdp committed 234 days ago

8f9ce9e9

[pallas:mosaic] Pulled `aval_to_ir_type` into `LoweringRuleContext`

superbobry committed 234 days ago

c2b10e54

Include debug key reuse config in the jit key

apaszke committed 234 days ago

8b9a9d53

[mosaic] Verify that `tpu.memref_slice` always produces tile-aligned slices

superbobry committed 234 days ago

4a14e0c2

[jaxlib] Add `__eq__` and `__hash__` to `PyDevice`.

danielsuo committed 234 days ago

279a9269

Fix reference counting bug in safe_map and for_each on error path.

hawkinsp committed 234 days ago

b1939f43

[Pallas:SC] Adds support for bool reduces, and min reduce.

brianwa84 committed 234 days ago

042c3abf

Older