ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 213 days ago

e8f4457a

Merge pull request #29570 from carlosgmartin:gitignore_build_editable

Google-ML-Automation committed 213 days ago

fdac4e68

Merge pull request #30077 from jakevdp:cleanup-unused

Google-ML-Automation committed 213 days ago

12847df7

Simplify `as_manual_mesh` now that shard_map only tracks manual axes and not `auto` as it did previously.

yashk2810 committed 213 days ago

ac9c4c42

[Pallas TPU] Enable constraining memory spaces of Refs and closing over them in core_map.

sharadmv committed 213 days ago

1447ef05

jax/BUILD: switch to pytype_strict_library everywhere

Jake VanderPlas committed 213 days ago

480b4992

remove unused internal function

jakevdp committed 213 days ago

56870252

Finalize deprecation of jax.lib.xla_client.DeviceAssignment

Jake VanderPlas committed 214 days ago

9e587572

[Pallas TPU] Add lookahead for input buffers to emit_pipeline.

justinjfu committed 214 days ago

6183b127

[Mosaic] Add tpu.wait_indirect_dma

naummo committed 214 days ago

e2d61be6

Merge pull request #30037 from skye:jas_link

Google-ML-Automation committed 214 days ago

10b9fe00

[Pallas] Make CommsEffect lowerable

sharadmv committed 214 days ago

20d6c1dd

Add a few tests to check lowering cache hits

yashk2810 committed 214 days ago

a46be09c

Merge pull request #28839 from carlosgmartin:print_environment_info_environment_variables_xla

Google-ML-Automation committed 214 days ago

77acfb3e

Merge pull request #29997 from carlosgmartin:tree_reduce_associative

Google-ML-Automation committed 214 days ago

966bf01f

[Pallas] Properly handle consts in while_loop fusible_dtype rule.

justinjfu committed 214 days ago

787c1173

Add jax.tree.reduce_associative.

carlosgmartin committed 214 days ago

7343e14d

Merge pull request #30071 from carlosgmartin:remove_print_from_tree_util_test_testStringRepresentation

Google-ML-Automation committed 214 days ago

6d6ef7d6

Remove print statement from tree_util_test.TreeTest.testStringRepresentation.

carlosgmartin committed 214 days ago

d7ebdb98

Merge pull request #30060 from jakevdp:changelog-update

Google-ML-Automation committed 214 days ago

251ac072

Reverts e562cddc3421f3ee10a031ae21c48ed85f3edbb2

emilyfertig committed 214 days ago

f4c1715a

[Pallas TPU] Fix bug where we wouldn't pass in a mesh context during lowering because we'd fail to detect comms in the Jaxpr

sharadmv committed 214 days ago

9f9c2d05

Merge pull request #30034 from mwhittaker:heartbeat_in_initialize

Google-ML-Automation committed 214 days ago

ae80a5dc

[Mosaic GPU] Enable `mypy` and fix type errors for `dialect_lowering.py`.

bchetioui committed 214 days ago

d5b5d522

[Mosaic GPU] Remove a few stale `jaxlib` guards.

bchetioui committed 214 days ago

9c26bbe7

Fix checking presence of `in_devices` in `_get_global_axis_size`.

cantonios committed 214 days ago

e38c1285

[Mosaic GPU] Add `optimization_barrier` to the new equational layout inference.

bchetioui committed 214 days ago

2d3b24ce

[Mosaic GPU] Add support for reductions across a subset of warps

apaszke committed 214 days ago

896b6e83

[Mosaic GPU] Make sure that TMEMLayout.canonicalize returns a TMEMLayout

apaszke committed 214 days ago

0c087ea6

[Mosaic GPU] Add support for elementwise ops in the new equational layout inference.

bchetioui committed 214 days ago

798320f0

Older