ROCm/jax

Pull Requests Commits

[sharding_in_types] Add `core.ShardingTypeError` as a new Exception that are sharding-in-types specific errors should raise.

yashk2810 committed 330 days ago

c7d6b653

[sharding_in_types] Fix a dynamic_slice bug where in the transpose, `DUS`'s operand was not sharded properly

yashk2810 committed 330 days ago

a8fb0e01

Merge pull request #27034 from vfdev-5:enable-mt-jax-context

Google-ML-Automation committed 330 days ago

90b8820b

Enable multi-threading in Jax Context with shared thread pool

vfdev-5 committed 330 days ago

55b55e6b

[Mosaic GPU] Add preliminary TMEM allocation support for Pallas/Mosaic GPU.

justinjfu committed 330 days ago

d0b71fa1

Merge pull request #26387 from Rifur13:sharding

Google-ML-Automation committed 330 days ago

80784a5b

Allow `tuple(arrays)` as an input to `make_array_from_single_device_arrays`. Fixes https://github.com/jax-ml/jax/issues/27303

yashk2810 committed 330 days ago

ea7fa29b

Adding sharding support to dynamic masks

Rifur13 committed 330 days ago

412f1d35

[Mosaic GPU] Skip Mosaic GPU tests if jax_pallas_use_mosaic_gpu flag is not set.

justinjfu committed 330 days ago

59e480db

Merge pull request #27107 from pearu:pearu/betainc

Google-ML-Automation committed 330 days ago

5745ff52

Fix process_allgather of global jax.Arrays with shardy

yashk2810 committed 330 days ago

1ec05853

Reverts 2562da7026ccd930e5f0972598c7d5479175b787

Google-ML-Automation committed 330 days ago

dad1b41f

[JAX Shardy] #sdy Unskip another test that is now passing

tomnatan30 committed 330 days ago

8bbd738d

[JAX Shardy] Unskip stream annotation test when shardy is enabled, since the underlying issue is now resolved.

tomnatan30 committed 330 days ago

c098b363

Update XLA dependency to use revision

Google-ML-Automation committed 330 days ago

7fa7db7a

Reorder C++ imports.

hawkinsp committed 330 days ago

2c90fe2d

[Mosaic GPU] Add support for slicing tiled refs with (tile aligned) dynamic base offsets

apaszke committed 330 days ago

4d6f15f2

Update XLA dependency to use revision

Google-ML-Automation committed 330 days ago

1c8e60e6

[XLA:GPU][Triton] Remove sparsity code.

chsigg committed 331 days ago

84cc397b

Remove XLA FFI GPU callback handler.

danielsuo committed 331 days ago

f1298ae7

Deprecate the jaxlib.hlo_helpers submodule.

dfm committed 331 days ago

e2b6859e

[mosaic_gpu] Don't time the warmup step in cupti profiler.

chr1sj0nes committed 331 days ago

18326abe

[Mosaic GPU] Introduce an optimization barrier op.

bchetioui committed 331 days ago

2d43fb47

[Mosaic:TPU] Add overload to ComputeTileStrides that just takes a shape.

tlongeri committed 331 days ago

6e204171

[mosaic_gpu] Make cupti finalization optional.

chr1sj0nes committed 331 days ago

509c6589

Merge pull request #27273 from carlosgmartin:fix_categorical_docs

Google-ML-Automation committed 331 days ago

761bd42c

Merge pull request #27274 from yliu120:new_fix_annotation

Google-ML-Automation committed 331 days ago

bb274f13

[mosaic_gpu] Check for dropped activity records in cupti profiler.

chr1sj0nes committed 331 days ago

58ba4106

Reverts e0c093314d8d9a6f68953f0c340c1b01d50ce386

Google-ML-Automation committed 331 days ago

4da751a9

Remove ; in code blocks of `thinking_in_jax.md`

Google-ML-Automation committed 331 days ago

e0c09331

Older