ROCm/jax

Pull Requests Commits

Merge pull request #26281 from jakevdp:lax-docs

Google-ML-Automation committed 1 year ago

fdcc04c3

Merge pull request #26278 from jax-ml:dependabot/github_actions/actions/setup-python-5.4.0

Google-ML-Automation committed 1 year ago

023f06f9

Bump actions/setup-python from 5.3.0 to 5.4.0

dependabot[bot] committed 1 year ago

Verified 333c4e7a

Merge pull request #26310 from dfm:use-uv-in-ci

Google-ML-Automation committed 1 year ago

694572bc

Add get/set methods to TransformedRef

sharadmv committed 1 year ago

782215b0

Merge pull request #26078 from gnecula:debug_info_jaxpr

Google-ML-Automation committed 1 year ago

414449e1

fix: simplify `logcdf` implementation by removing unnecessary argument promotion

Qazalbash committed 1 year ago

Verified 8561f90f

Merge pull request #26307 from dfm:kill-ffi-cache

Google-ML-Automation committed 1 year ago

b1b88a36

jax.lax: improve docs for exp & log functions

jakevdp committed 1 year ago

f180353d

Merge pull request #26302 from vfdev-5:readd-missed-sed-part-in-tsan

Google-ML-Automation committed 1 year ago

09ee37a4

Merge pull request #26300 from vfdev-5:patch-2

Google-ML-Automation committed 1 year ago

2b944442

Use uv for dependency resolution on CI.

dfm committed 1 year ago

5db5e0d5

fix binomial sample test on TPU

froystig committed 1 year ago

58a202ad

[Mosaic GPU] Add simple Blackwell matmul example

andportnoy committed 1 year ago

aff2cba8

Remove the FFI examples workflow cache.

dfm committed 1 year ago

42b8eb6c

Remove old backward compatibility mode for old PRGN custom call on GPU

gnecula committed 1 year ago

9f797990

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

6738986c

Readded missing part to use locally built numpy in the build in tsan.yml

vfdev-5 committed 1 year ago

fc408d20

Set cron to 05:00 UTC in tsan.yaml

vfdev-5 committed 1 year ago

Verified a026de23

[pallas_mgpu] Add a test for emit_pipeline with wgmma.

petebu committed 1 year ago

c7d535d3

[better_errors] Add debug info to more Jaxprs and WrappedFun (step 1)

gnecula committed 1 year ago

d12aead6

[Pallas] Support promise_in_bounds mode in jnp.take_along_axis.

bythew3i committed 1 year ago

124e1239

Strip all info except for shape and dtype while matching input and output avals for donation. For avals that are sharded, we match on shard_shape.

yashk2810 committed 1 year ago

654a2f6e

fix: update `logcdf` implementation to use `sf`

Qazalbash committed 1 year ago

9a324bfb

Merge pull request #26289 from reikdas:reikdas/docfix

Google-ML-Automation committed 1 year ago

6281b860

[sharding_in_types] Add a canonicalize_value step before dispatching `bind` so that we can insert `mesh_cast`s under the following conditions:

yashk2810 committed 1 year ago

bc1a7066

Propagate source ranges in location.

jpienaar committed 1 year ago

60d3836f

Merge pull request #26290 from mattjj:linearize-name-stack-fixes-2

Google-ML-Automation committed 1 year ago

363f1e69

`try_matching_out_with_in_spec_for_all_auto` should only work for NamedSharding

yashk2810 committed 1 year ago

5e7e6911

[direct-linearize] fix name stack tests

mattjj committed 1 year ago

8f967c51

Newer Older