ROCm/jax

Pull Requests Commits

Remove unused `return wrapper` in annotate_function that creates a self reference cycle loop in python.

Google-ML-Automation committed 270 days ago

73ecf0bb

[pallas:mosaic_gpu] `emit_pipeline*` now allows the grid to be dynamic

superbobry committed 270 days ago

5f33280d

Removed `eager_pmap` config option

superbobry committed 270 days ago

12811f08

[Mosaic GPU] Add warpgroup lowering for `RunState` in Pallas.

dimitar-asenov committed 270 days ago

c4cc94a1

[Mosaic GPU] Refactor and generalize code in `optimization_barrier`.

dimitar-asenov committed 270 days ago

d12cbffd

Removed redundant `pass`es

superbobry committed 270 days ago

af072feb

Removed `jax._src.raise_to_shaped`

superbobry committed 270 days ago

8ed59d8b

Merge pull request #27776 from gnecula:export_keys

Google-ML-Automation committed 270 days ago

c2eaedfe

Clarify jax.make_jaxpr docstring

j-towns committed 270 days ago

bc11a631

[Mosaic GPU] Add support for replicated warp_dim parsing and a dedicated test for parsing all canonical layouts.

dimitar-asenov committed 270 days ago

19fcae12

[export] Add backwards compatibility test for annotate_device_placement.

gnecula committed 270 days ago

51dbcd4d

Address previous FP8-related TODOs in jaxlib/XLA.

apivovarov committed 270 days ago

bb515aa7

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

hawkinsp committed 270 days ago

86de4783

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

dfm committed 270 days ago

31589960

Merge pull request #27814 from ZacCranko:harden-cache

Google-ML-Automation committed 270 days ago

4bae9cda

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

Google-ML-Automation committed 270 days ago

9e036865

harden cache against jaxlib ver

ZacCranko committed 270 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 270 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 270 days ago

0a72e856

Removed `data_dependent_tracing_fallback` config option

superbobry committed 270 days ago

2944e3b2

Merge pull request #27716 from jakevdp:jax-array

Google-ML-Automation committed 270 days ago

34205468

Reverts 006a6a63feb64bf9984526030ba008186d69d2b4

Google-ML-Automation committed 270 days ago

48a9ad07

[shard-map] canonicalize rep=None to be rep={all possible axes}

mattjj committed 270 days ago

3a3c1450

Relax jax dependency constraints to be able to install RC wheels

nitins17 committed 270 days ago

64e4bf26

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

justinjfu committed 270 days ago

b18dc1df

jnp.linalg: add symmetrize_input argument & docs

jakevdp committed 270 days ago

96e63eae

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

hawkinsp committed 270 days ago

db11efab

jax.numpy: support __jax_array__ in remaining APIs

jakevdp committed 270 days ago

d3cfff05

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

Google-ML-Automation committed 270 days ago

74917ce5

[shard-map] add while_map rep rule

mattjj committed 270 days ago

9a3e94de

Newer Older