ROCm/jax

Pull Requests Commits

Merge pull request #27273 from carlosgmartin:fix_categorical_docs

Google-ML-Automation committed 1 year ago

761bd42c

Merge pull request #27274 from yliu120:new_fix_annotation

Google-ML-Automation committed 1 year ago

bb274f13

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

chr1sj0nes committed 1 year ago

58ba4106

Reverts e0c093314d8d9a6f68953f0c340c1b01d50ce386

Google-ML-Automation committed 1 year ago

4da751a9

Remove ; in code blocks of `thinking_in_jax.md`

Google-ML-Automation committed 1 year ago

e0c09331

Fixes the stream annotation compute on box.

yliu120 committed 1 year ago

258ed1b0

Fix doc for random.categorical replace argument.

carlosgmartin committed 1 year ago

b5c467e6

Expose profiler_data submodule from XLA to Jaxlib.

Google-ML-Automation committed 1 year ago

2562da70

Merge pull request #27254 from jburnim:jburnim_pallas_interpret_mode5

Google-ML-Automation committed 1 year ago

fc97b0cb

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

ezhulenev committed 1 year ago

47dde87b

Fix `lax_autodiff_test` on v5p

Google-ML-Automation committed 1 year ago

f7471125

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

Google-ML-Automation committed 1 year ago

16dc0ad1

Add a presubmit check to test against oldest supported numpy

nitins17 committed 1 year ago

29e90a30

Remove code to support jaxlib < 0.5.3.

hawkinsp committed 1 year ago

362fb7ae

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

yashk2810 committed 1 year ago

4489303d

Merge pull request #27260 from hawkinsp:postrelease

Google-ML-Automation committed 1 year ago

04454b65

Merge pull request #27255 from pearu:pearu/assertAllClose

Google-ML-Automation committed 1 year ago

afdbcd02

Merge pull request #27258 from jakevdp:fix-lint

Google-ML-Automation committed 1 year ago

cf21f733

Update version numbers after JAX 0.5.3 release.

hawkinsp committed 1 year ago

9d534ad2

Merge branch 'release/0.5.3' into main

hawkinsp committed 1 year ago

5d4de831

Fix lint error on main

jakevdp committed 1 year ago

7a67c9bd

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

pearu committed 1 year ago

5a5415bc

Add sliding window support to the ragged paged attention.

Google-ML-Automation committed 1 year ago

84ec21e0

Move sparse op GPU lowerings from jaxlib into JAX.

dfm committed 1 year ago

918192fd

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

superbobry committed 1 year ago

b456855c

Support error checking in auto mode

ayaka14732 committed 1 year ago

85e78840

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

hawkinsp committed 1 year ago

ee74c289

[JAX] Move py_client_gpu into JAX.

hawkinsp committed 1 year ago

dd93eeae

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

yashk2810 committed 1 year ago

dde861af

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

jburnim committed 1 year ago

af5b2efd

Older