ROCm/jax

Pull Requests Commits

Merge pull request #24602 from rdyro:rdyro-decode-attention-mask

Google-ML-Automation committed 1 year ago

37af1002

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

dc33a280

Remove some obsolete deprecation registrations

Jake VanderPlas committed 1 year ago

4e45a6d9

Merge pull request #24733 from jakevdp:fix-quantile

Google-ML-Automation committed 1 year ago

36017a4a

[Pallas:TPU] Fix some stale/wrong skip conditions.

WindQAQ committed 1 year ago

b6f5c95a

[Mosaic] Verify that tpu.sem_wait semaphore rank is zero

naummo committed 1 year ago

3df204a4

Merge pull request #24723 from jakevdp:beta-dep

Google-ML-Automation committed 1 year ago

cbaafbbe

Fix a bug in jax.scipy.stats.rankdata leading to breakage with shape polymorphism.

Google-ML-Automation committed 1 year ago

542cb2e5

scipy.special.beta: remove deprecated x and y parameters

jakevdp committed 1 year ago

d698da61

Merge pull request #24722 from jakevdp:tracer-hash

Google-ML-Automation committed 1 year ago

0fcb2f39

Add support for tpu v5e to `jax.make_mesh`

yashk2810 committed 1 year ago

b8c263a5

Adding start index and kv_seq_len to decode kernel

rdyro committed 1 year ago

d62510bf

Fix debug_nans false positive in jnp.quantile

jakevdp committed 1 year ago

44c6883c

Merge pull request #24729 from rdyro:logging-fix

Google-ML-Automation committed 1 year ago

4b4fb9da

Include mpmath as a bazel dependency of lax_test.

hawkinsp committed 1 year ago

ea1e8795

rdyro committed 1 year ago

b60d0ab8

Merge pull request #24704 from andportnoy:patch-3

Google-ML-Automation committed 1 year ago

563ecdf2

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

21f4b085

Merge pull request #24717 from jakevdp:fix-rankdata

Google-ML-Automation committed 1 year ago

939b41f5

Merge pull request #24710 from rajasekharporeddy:typos

Google-ML-Automation committed 1 year ago

c1af808c

Merge pull request #23468 from rdyro:rdyro-add-logging-env

Google-ML-Automation committed 1 year ago

497a5a35

Merge pull request #23969 from gspschmid:gschmid/dist-compression

Google-ML-Automation committed 1 year ago

a70e7993

Adding `JAX_LOGGING_LEVEL` configuration option

rdyro committed 1 year ago

04f2ef9e

Make Tracers non-hashable

jakevdp committed 1 year ago

095bb0e7

[jax.distributed] Enable grpc channel compression

gspschmid committed 1 year ago

7bdb2bf9

Reverts a913fbf2fddc5b8c1b6c85b159d0eeb1bf65d461

hawkinsp committed 1 year ago

0e8acff5

Reverts f281c6f46475270a57a02416469226315377592c

dougalm committed 1 year ago

478b750c

Improve efficiency of jax.scipy.stats.rankdata

jakevdp committed 1 year ago

5f90f63d

[mosaic_gpu] Check the return code of `gpuEventCreate` and `gpuEventDestroy`

superbobry committed 1 year ago

34b4787e

[Mosaic GPU] Ensure that the dialect module can be loaded successfully.

bchetioui committed 1 year ago

63e59c5f

Older