ROCm/jax

Pull Requests Commits

Merge pull request #27273 from carlosgmartin:fix_categorical_docs

Google-ML-Automation committed 303 days ago

761bd42c

Merge pull request #27274 from yliu120:new_fix_annotation

Google-ML-Automation committed 303 days ago

bb274f13

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

chr1sj0nes committed 303 days ago

58ba4106

Reverts e0c093314d8d9a6f68953f0c340c1b01d50ce386

Google-ML-Automation committed 303 days ago

4da751a9

Remove ; in code blocks of `thinking_in_jax.md`

Google-ML-Automation committed 303 days ago

e0c09331

Fixes the stream annotation compute on box.

yliu120 committed 303 days ago

258ed1b0

Fix doc for random.categorical replace argument.

carlosgmartin committed 303 days ago

b5c467e6

Expose profiler_data submodule from XLA to Jaxlib.

Google-ML-Automation committed 303 days ago

2562da70

Merge pull request #27254 from jburnim:jburnim_pallas_interpret_mode5

Google-ML-Automation committed 303 days ago

fc97b0cb

Use np.ones to avoid signed integer overflow at run time

ezhulenev committed 303 days ago

47dde87b

Fix `lax_autodiff_test` on v5p

Google-ML-Automation committed 303 days ago

f7471125

Add `jax_source_package` macros and target to generate a source package `.tar.gz`.

Google-ML-Automation committed 303 days ago

16dc0ad1

Add a presubmit check to test against oldest supported numpy

nitins17 committed 303 days ago

29e90a30

Remove code to support jaxlib < 0.5.3.

hawkinsp committed 303 days ago

362fb7ae

Delete `ParsedPartitionSpec` and `preprocess` function and do a couple more cleanups

yashk2810 committed 303 days ago

4489303d

Merge pull request #27260 from hawkinsp:postrelease

Google-ML-Automation committed 303 days ago

04454b65

Merge pull request #27255 from pearu:pearu/assertAllClose

Google-ML-Automation committed 303 days ago

afdbcd02

Merge pull request #27258 from jakevdp:fix-lint

Google-ML-Automation committed 303 days ago

cf21f733

Update version numbers after JAX 0.5.3 release.

hawkinsp committed 304 days ago

9d534ad2

Merge branch 'release/0.5.3' into main

hawkinsp committed 304 days ago

5d4de831

Fix lint error on main

jakevdp committed 304 days ago

7a67c9bd

Rename arguments x, y of assertAllClose and friends to actual, expected.

pearu committed 304 days ago

5a5415bc

Add sliding window support to the ragged paged attention.

Google-ML-Automation committed 304 days ago

84ec21e0

Move sparse op GPU lowerings from jaxlib into JAX.

dfm committed 304 days ago

918192fd

[pallas:mosaic_gpu] Added support for accessing cluster ID via `lax.axis_index`

superbobry committed 304 days ago

b456855c

Support error checking in auto mode

ayaka14732 committed 304 days ago

85e78840

Move //jaxlib:handle_pool to //jaxlib/gpu:handle_pool.

hawkinsp committed 304 days ago

ee74c289

[JAX] Move py_client_gpu into JAX.

hawkinsp committed 304 days ago

dd93eeae

Remove the jax Array migration guide from the TOC tree but keep the doc around

yashk2810 committed 304 days ago

dde861af

Fix input_output_aliases for non-HBM kernel args in TPU interpret mode.

jburnim committed 304 days ago

af5b2efd

Older