ROCm/jax

Pull Requests Commits

Address review comments

alekstheod committed 58 days ago

85b3ffbd

Use rbe certs from rocAutomation

alekstheod committed 58 days ago

72b938b3

alekstheod committed 58 days ago

66893d83

Overwrite rocm CI job with a job that executes the tests

alekstheod committed 58 days ago

c21684ea

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

Google-ML-Automation committed 58 days ago

aee33099

Add memory_space to typematch check

yashk2810 committed 59 days ago

540df654

Merge pull request #35095 from carlosgmartin:lax_reduce_doc_init_value_scalar

Google-ML-Automation committed 59 days ago

4a374b41

Merge pull request #35153 from davisyoshida:hipspec

Google-ML-Automation committed 59 days ago

7ecbc030

Export HipSpec from hijax

davisyoshida committed 59 days ago

fc641811

Merge pull request #35139 from jakevdp:fix-lint

Google-ML-Automation committed 59 days ago

ab7aecd7

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

Google-ML-Automation committed 59 days ago

80ebaef8

Restore scan length check

jakevdp committed 59 days ago

1eed1171

Reverts 4a14e0c2cce9ff41cbca271b429d36df8215e608

nvgrw committed 59 days ago

5abf880a

Fix jupytext lint at head

jakevdp committed 59 days ago

9b317766

Merge pull request #35137 from mattjj:issue34758

Google-ML-Automation committed 59 days ago

81640be3

Merge pull request #35136 from mattjj:hi-vmap

Google-ML-Automation committed 59 days ago

1c1ad14e

[hijax] test vmap + hijax types + hijax primitive

mattjj committed 59 days ago

81ebe31f

[hijax] add cur_qdd to BatchTrace

mattjj committed 59 days ago

b456d935

Merge pull request #35083 from mattjj:hi-shmap

Google-ML-Automation committed 59 days ago

0ea5ee02

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

mattjj committed 59 days ago

cf879e1b

Integrate LLVM at llvm/llvm-project@c6bb1afbc508

Google-ML-Automation committed 59 days ago

6a0d96d3

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

superbobry committed 59 days ago

70692c6f

Disable an einshape test on TPU 7x that is failing

hawkinsp committed 59 days ago

696a10ae

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

Google-ML-Automation committed 59 days ago

a10ee0ae

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

dimitar-asenov committed 59 days ago

e0b75368

[pyrefly] more typing fixes

jakevdp committed 59 days ago

8f9ce9e9

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

superbobry committed 59 days ago

c2b10e54

Include debug key reuse config in the jit key

apaszke committed 59 days ago

8b9a9d53

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

superbobry committed 59 days ago

4a14e0c2

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

danielsuo committed 59 days ago

279a9269

Older