ROCm/jax

Pull Requests Commits

Silence buildifier string list warnings

Jake VanderPlas committed 1 year ago

bf45baf5

print `pjit_p` as `"jit"` (in jaxpr, etc.)

froystig committed 1 year ago

d585aecc

[Pallas] Pass a `metadata` dict to mosaic

sharadmv committed 1 year ago

3f5ed191

Merge pull request #29864 from mfschubert:main

Google-ML-Automation committed 1 year ago

518975e6

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

4716d0b3

Don't calculate `scan_xs` and `scan_ys` when batch size is greater than len(input) i.e. num_batches == 0. Fixes https://github.com/jax-ml/jax/issues/29867

yashk2810 committed 1 year ago

fe1cef42

[doc] Clarify Profiling Docs for XProf and Tensorboard integration

Matt-Hurd committed 1 year ago

712fa816

Add memory_space to state.AbstractRef and use it instead of AbstractMemoryRef in Pallas

sharadmv committed 1 year ago

7cf7c967

Change parameters of `py_import` in accordance with https://github.com/openxla/xla/pull/28179.

Google-ML-Automation committed 1 year ago

3aaa134a

Move `jax._src.debugger` to its own build rule

Jake VanderPlas committed 1 year ago

6a29f08f

Move `jax._src.blocked_sampler` to its own build rule

Jake VanderPlas committed 1 year ago

547c894b

Merge pull request #29868 from mattjj:fix-remat-computation-annotation

Google-ML-Automation committed 1 year ago

d8e58e38

[remat] fix 'rematted_computation' annotation

mattjj committed 1 year ago

6a0eceb2

[Mosaic] Add DMA source and target shape verification

naummo committed 1 year ago

03a8df1e

Update gpu_memory_allocation.rst

mfschubert committed 1 year ago

Verified 5a320c0e

Show stdout and stderr on test failure.

mwhittaker committed 1 year ago

1becc8d8

Move jax._src.random to its own build target

Jake VanderPlas committed 1 year ago

f383e1cc

Fix an edge-case where the reshape has a split with extra 1s at the end in both orig_shape and new_sizes. Fixes https://github.com/jax-ml/jax/issues/29859

yashk2810 committed 1 year ago

71a132e0

Change ValueError to NotImplementedError when setting partitioned=True on callbacks for non-CPU/GPU devices.

danielsuo committed 1 year ago

ee264800

[CI] Relax the timeout constraint on pytest_tpu tests

MichaelHudgins committed 1 year ago

d2ed0244

If the context mesh is fully explicit and val mesh is empty, then set the out_sharding to be fully replicated. Fixes https://github.com/jax-ml/jax/issues/29837

yashk2810 committed 1 year ago

3816b4f9

[Mosaic GPU] Fix another small bug in FragmentedArray.transfer_tiled

apaszke committed 1 year ago

191b8946

[JAX] Remove lax.infeed and lax.outfeed from JAX's public APIs.

hawkinsp committed 1 year ago

239a4a0e

[Pallas:TPU] Bump the too low minimum libtpu version for a test

apaszke committed 1 year ago

59ba4939

[Mosaic GPU] Add an initial skeleton for an equation system for layout/transform inference.

bchetioui committed 1 year ago

d08c35fb

[Mosaic GPU] Fix a minor bug in FragmentedArray.transfer_tiled

apaszke committed 1 year ago

6d188eee

Merge pull request #29539 from gnecula:debug_info1

Google-ML-Automation committed 1 year ago

fcc0c580

[debug_info] Improve debug info for loop constructs.

gnecula committed 1 year ago

16ffcf39

[Mosaic GPU] Fix failures in B200 CI

apaszke committed 1 year ago

0e76f73a

[Mosaic GPU] Fix bug in `Tiling.remove_dimension`.

bchetioui committed 1 year ago

ee832c1e

Older