ROCm/jax

Pull Requests Commits

Address previous FP8-related TODOs in jaxlib/XLA.

apivovarov committed 1 year ago

bb515aa7

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

hawkinsp committed 1 year ago

86de4783

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

dfm committed 1 year ago

31589960

Merge pull request #27814 from ZacCranko:harden-cache

Google-ML-Automation committed 1 year ago

4bae9cda

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

Google-ML-Automation committed 1 year ago

9e036865

harden cache against jaxlib ver

ZacCranko committed 1 year 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 1 year ago

84e04fe6

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

yashk2810 committed 1 year ago

0a72e856

Removed `data_dependent_tracing_fallback` config option

superbobry committed 1 year ago

2944e3b2

Merge pull request #27716 from jakevdp:jax-array

Google-ML-Automation committed 1 year ago

34205468

Reverts 006a6a63feb64bf9984526030ba008186d69d2b4

Google-ML-Automation committed 1 year ago

48a9ad07

Relax jax dependency constraints to be able to install RC wheels

nitins17 committed 1 year ago

64e4bf26

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

justinjfu committed 1 year ago

b18dc1df

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

hawkinsp committed 1 year ago

db11efab

jax.numpy: support __jax_array__ in remaining APIs

jakevdp committed 1 year ago

d3cfff05

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

Google-ML-Automation committed 1 year ago

74917ce5

Add jaxlib_extension_version guard against explicit copying

pschuh committed 1 year ago

b6e4b938

[mgpu] Allow bf16 printing

cperivol committed 1 year ago

e1b05728

Merge pull request #27797 from mattjj:shmap-fix

Google-ML-Automation committed 1 year ago

0d067313

[CI] Temporarily disable TPU v6 due to runner issues

MichaelHudgins committed 1 year ago

522add2c

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

23b63cd5

Apply output forwarding in lin rule for pjit.

dfm committed 1 year ago

dc00f9bd

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

mattjj committed 1 year ago

05ca0233

Merge pull request #26853 from jeffcarp:scalar-event

Google-ML-Automation committed 1 year ago

f23dd642

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

bythew3i committed 1 year ago

e1e37f8d

[Pallas Fuser] Add output_fusion_mask support

sharadmv committed 1 year ago

fcf5115f

Bump medyagh/setup-minikube from 0.0.18 to 0.0.19

dependabot[bot] committed 1 year ago

Verified 7239487c

Add int4, uint4 to test_util.suppported_types

jburnim committed 1 year ago

855829e1

Merge pull request #27735 from dfm:lin-fwd

Google-ML-Automation committed 1 year ago

5581e7d0

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

superbobry committed 1 year ago

51c224c4

Older