ROCm/jax

Pull Requests Commits

Address previous FP8-related TODOs in jaxlib/XLA.

apivovarov committed 329 days ago

bb515aa7

Remove unused function jax._src.interpreters.mlir.xla_computation_to_mlir_module.

hawkinsp committed 329 days ago

86de4783

Migrate custom_call filecheck to use internal custom_call since the external one is deprecated.

dfm committed 329 days ago

31589960

Merge pull request #27814 from ZacCranko:harden-cache

Google-ML-Automation committed 329 days ago

4bae9cda

Merge pull request #27793 from dfm:lin-out-fwd

Google-ML-Automation committed 329 days ago

9e036865

harden cache against jaxlib ver

ZacCranko committed 329 days ago

ca6e470d

Add custom pretty print rule for the unary ops with accuracy s.t. accuracy is not printed if it's None.

hanrach9 committed 329 days ago

84e04fe6

Add **experimental** `with_dll_constraint` API. This is for cases when the users wants to let SPMD decide the sharding.

yashk2810 committed 329 days ago

0a72e856

Removed `data_dependent_tracing_fallback` config option

superbobry committed 329 days ago

2944e3b2

Merge pull request #27716 from jakevdp:jax-array

Google-ML-Automation committed 329 days ago

34205468

Reverts 006a6a63feb64bf9984526030ba008186d69d2b4

Google-ML-Automation committed 329 days ago

48a9ad07

Relax jax dependency constraints to be able to install RC wheels

nitins17 committed 329 days ago

64e4bf26

[Mosaic GPU] Add scaffolding for a new lowering "axis" (UserThreadSemantics), in addition to the existing ThreadSemantics (renamed to LoweringSemantics).

justinjfu committed 329 days ago

b18dc1df

Migrate jaxlib to use a single common .so file for all C++ dependencies.

hawkinsp committed 329 days ago

db11efab

jax.numpy: support __jax_array__ in remaining APIs

jakevdp committed 329 days ago

d3cfff05

Merge pull request #27794 from jax-ml:dependabot/github_actions/medyagh/setup-minikube-0.0.19

Google-ML-Automation committed 329 days ago

74917ce5

Add jaxlib_extension_version guard against explicit copying

pschuh committed 329 days ago

b6e4b938

[mgpu] Allow bf16 printing

cperivol committed 329 days ago

e1b05728

Merge pull request #27797 from mattjj:shmap-fix

Google-ML-Automation committed 329 days ago

0d067313

[CI] Temporarily disable TPU v6 due to runner issues

MichaelHudgins committed 329 days ago

522add2c

Update XLA dependency to use revision

Google-ML-Automation committed 329 days ago

23b63cd5

Apply output forwarding in lin rule for pjit.

dfm committed 329 days ago

dc00f9bd

[shard-map] in eager shmap, handle all rep rule output cases

mattjj committed 329 days ago

05ca0233

Merge pull request #26853 from jeffcarp:scalar-event

Google-ML-Automation committed 329 days ago

f23dd642

[Mosaic TPU] FWD compatibility needs to keep previous version at least one month.

bythew3i committed 329 days ago

e1e37f8d

[Pallas Fuser] Add output_fusion_mask support

sharadmv committed 329 days ago

fcf5115f

Bump medyagh/setup-minikube from 0.0.18 to 0.0.19

dependabot[bot] committed 329 days ago

Verified 7239487c

Add int4, uint4 to test_util.suppported_types

jburnim committed 329 days ago

855829e1

Merge pull request #27735 from dfm:lin-fwd

Google-ML-Automation committed 329 days ago

5581e7d0

Removed deprecated `jax.core.{full_lower,jaxpr_as_fun,lattice_join}`

superbobry committed 329 days ago

51c224c4

Older