ROCm/jax

Pull Requests Commits

Update rules_python.patch to support Python 3.13.0 and update python 3.13 packages in JAX.

Google-ML-Automation committed 1 year ago

53668b88

Merge pull request #24172 from shuhand0:dev/shuhan/adopt_jaxlib0.4.34

Google-ML-Automation committed 1 year ago

bcb0f646

Merge pull request #24214 from dfm:api-docs-make-mesh

Google-ML-Automation committed 1 year ago

4025b27a

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

d36afe4f

Add jax.make_mesh to API docs.

dfm committed 1 year ago

1f0a04a4

Merge pull request #24167 from rajasekharporeddy:testbranch1

Google-ML-Automation committed 1 year ago

db71965c

Merge pull request #24210 from jakevdp:nan-to-num-doc

Google-ML-Automation committed 1 year ago

c2deae8a

[Pallas TPU] Fix comparison lowering for unsigned integers

ayaka14732 committed 1 year ago

77613f21

#sdy use updated axis context when skipping `ManualComputationOp` creation.

bartchr808 committed 1 year ago

875f44c6

Better documentation for jnp.nan_to_num

jakevdp committed 1 year ago

2f798902

Better docs for jnp.left_shift

rajasekharporeddy committed 1 year ago

ed028be7

[Pallas] Add support for runtime checking of grid bounds using checkify.

justinjfu committed 1 year ago

9cf952a5

[JAX] Fix error message for matmul operand shape check.

Google-ML-Automation committed 1 year ago

9748e2ab

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

2f67710e

Error out in donation if:

yashk2810 committed 1 year ago

e5fa9656

Fix returning semaphores from pallas kernels with grids.

Google-ML-Automation committed 1 year ago

55153cc3

[mosaic_gpu] Memref reshape by means of folding/unfolding

cperivol committed 1 year ago

28b09342

Merge pull request #24189 from jakevdp:average-doc

Google-ML-Automation committed 1 year ago

54101771

cleanup: remove api_version from BUILD files

ericsalo committed 1 year ago

713e909b

Improve docs for jnp.average

jakevdp committed 1 year ago

0276514a

[Pallas/MGPU] Allow delaying the release of pipelined buffers

apaszke committed 1 year ago

25c1519a

Merge pull request #24174 from hawkinsp:asan

Google-ML-Automation committed 1 year ago

20cdddd7

Remove remaining implementations of jax.experimental.host_callback.call.

gnecula committed 1 year ago

023f2a78

Bump the shard_count of FFT tests to avoid timeouts

apaszke committed 1 year ago

7102c7ad

Move ASAN build to Python 3.13.0.

hawkinsp committed 1 year ago

667221ba

[Pallas] Simplify `OpsTest` by skipping 64-bit tests on 32-bit environments

ayaka14732 committed 1 year ago

6a958b90

Remove explicit_type argument from _nary_lower_hlo.

Google-ML-Automation committed 1 year ago

0854dc24

Use no_tracing config in `_create_pjit_jaxpr` to so that AOT path can also error if we re-trace.

yashk2810 committed 1 year ago

a9e9f97f

[pallas] Added `MemoryRef` and `run_scoped` to the API docs

superbobry committed 1 year ago

76d59380

Skip `test_ragged_copy_on_host` if `xla_extension_version` < 290

yashk2810 committed 1 year ago

ce2b4978

Newer Older