ROCm/jax

Pull Requests Commits

Move top-level _jax module contents into jax:: C++ namespace.

hawkinsp committed 154 days ago

d9b692f2

Prepare for removal of vector.ExtractElement and vector.InsertElement.

akuegel committed 154 days ago

c0bf76b3

[Mosaic GPU] Use a TiledLayout with vector_length=1 for 32-bit WGMMA accumulators

apaszke committed 154 days ago

7da7d4f4

[Pallas:MGPU] Add support for delay_release in the warpgroup specialized pipeline

apaszke committed 154 days ago

ad947e68

Merge pull request #30666 from gnecula:jaxpr_consts_shard_map

Google-ML-Automation committed 154 days ago

c15db488

[jaxpr consts] Fix const args handling in shard_map shardy lowering

gnecula committed 155 days ago

dbd853e6

[mutable-arrays] tweak title of array ref doc

mattjj committed 155 days ago

f227ccc4

Merge pull request #30675 from mattjj:array-ref-vjp

Google-ML-Automation committed 155 days ago

a2048d95

[mutable-arrays] scan autodiff fix, internal vjp-with-refs api prototype

mattjj committed 155 days ago

4020b146

Update XLA dependency to use revision

Google-ML-Automation committed 155 days ago

24f65283

[Pallas][Mosaic GPU] Support pytree in/out specs in the warp specialized pipeline.

Rifur13 committed 155 days ago

60087a4a

Raise an error if we see an invalid memory kind in device_put abstract_eval

yashk2810 committed 155 days ago

4a5e0e84

If timeout/transfer size flags are not passed, ensure that default values are used.

emilyfertig committed 155 days ago

7a288d4c

Migrate dependencies on jax experimental targets to `jax/experimental/BUILD`.

danielsuo committed 155 days ago

22cf005f

Merge pull request #30669 from mattjj:array-refs-docs

Google-ML-Automation committed 155 days ago

41297137

mattjj committed 155 days ago

4067a61a

[mutable-arrays] rename to ArrayRef, add jax.refs module

mattjj committed 155 days ago

1a946c0d

Remove the `RULES_PYTHON_REPO_DEBUG` flag from the docker container script

kanglant committed 155 days ago

7598e454

[Pallas/Fuser] Add pull rule for reduce_sum that fits whole dimension in VMEM.

sharadmv committed 155 days ago

6883f87c

Merge pull request #30662 from mattjj:array-ref-renaming

Google-ML-Automation committed 155 days ago

a2edb197

[Pallas] Remove unnecessary jit around core_map runnable

sharadmv committed 155 days ago

27a75b57

[mutable-arrays] rename to ArrayRef, add jax.refs module

mattjj committed 155 days ago

2c0f642b

Ignore the RuntimeWarning `invalid value encountered in isinf` in `testFrexp` method.

Google-ML-Automation committed 155 days ago

fd744266

Replace use_mesh with set_mesh everywhere in JAX

yashk2810 committed 155 days ago

f194b9aa

Rename xla::Traceback to jax::Traceback.

hawkinsp committed 155 days ago

4a8d78f7

Merge `use_mesh` and `set_mesh` into `set_mesh`. In other words, `set_mesh` can now be used as a global setter or a context manager.

yashk2810 committed 155 days ago

056cde49

[jax:experimental] Move names in experimental/__init__.py into their own targets.

danielsuo committed 155 days ago

8424ecad

[JAX] Remove deprecation of GSPMDSharding and PositionalSharding.

yashk2810 committed 155 days ago

26afd09c

More helpful error messages with big UUIDs.

pschuh committed 155 days ago

bb4a5b1c

Minor changes to sharding-in-types tutorial.

Ryan McKenna committed 155 days ago

3e742ff6

Newer Older