ROCm/jax

Pull Requests Commits

[Pallas] Pad input/outputs in interpret mode to fix errors in OOB memory accesses.

justinjfu committed 1 year ago

1e48adc6

Merge pull request #21119 from jakevdp:linalg-cond

a-googler committed 1 year ago

b8ed3466

Merge pull request #21048 from jakevdp:np-squeeze-doc

a-googler committed 1 year ago

6189a559

Async dispatch expensive computations on the JAX CPU backend.

yueshengys committed 1 year ago

9e7830df

Populate `propagated_out_mem_kinds` inside the branch where it's needed

yashk2810 committed 1 year ago

54d40727

jnp.linalg.cond: improve implementation & docs

jakevdp committed 1 year ago

1f6d9021

Merge pull request #21203 from gnecula:export_device_poly

a-googler committed 1 year ago

85e91c2b

Better documentation for several jax.numpy functions

jakevdp committed 1 year ago

7ceae95d

[export] Relax the check that exported modules are used with same number of devices as when exported

gnecula committed 1 year ago

98aead70

Remove type promotion for mixed fp8 matmuls.

elfiegg committed 1 year ago

43d19161

Merge pull request #21169 from justinjfu/splash_precision_fix

justinjfu committed 1 year ago

Verified e4f3b3ff

CI: update NumPy build version to 2.0.0rc2

Jake VanderPlas committed 1 year ago

35a512da

Reverts 49bd4d6f01d6cda00f9b1bdfbda156636baae928

cheshire committed 1 year ago

de14e3b3

Merge pull request #21191 from gnecula:export_simplify

a-googler committed 1 year ago

e66a234b

Merge pull request #21202 from superbobry:pallas

a-googler committed 1 year ago

54ca3d46

Merge pull request #20940 from piotrfilipiuk:changelist/623910451

a-googler committed 1 year ago

1fed7849

[export] Simplify export internals, prepare for integration with AOT APIs

gnecula committed 1 year ago

78d4d0a4

Guarded Pallas GPU import in tests/pallas/pallas_test.py

superbobry committed 1 year ago

8094d0d1

Ensured that all Pallas GPU tests depend on :pallas_gpu

superbobry committed 1 year ago

1c6855a4

Register TPU profiler plugin when get_topology_desc is called with tpu platform.

Jieying Luo committed 1 year ago

ba8480a2

[Mosaic GPU] Correct TileTransform.transform_index()

a-googler committed 1 year ago

af4bddb7

[jax:mosaic-gpu] Test type conversion for tiled fragment array

a-googler committed 1 year ago

dac5b754

[Mosaic GPU] Cleanup matmul example mainly by removing the m/n block loops (1 step).

a-googler committed 1 year ago

c1f184d6

Update XLA dependency to use revision

a-googler committed 1 year ago

3e3d916a

Update XLA dependency to use revision

a-googler committed 1 year ago

b4f21453

[Mosaic GPU] Prepare for writing warp-specialized kernels

apaszke committed 1 year ago

a527b719

Reverts 586568f4fe44cf9ad8b1bd022148a10c4b69f33a

hawkinsp committed 1 year ago

49bd4d6f

Implements Ragged Dot API

piotrfilipiuk committed 1 year ago

93dfe05a

Raise a runtime error when trying to convert the `jax.Array` wrapped by `jax.core.Token` to a numpy array, as it is an internal implementation detail and the buffer has XLA token shape.

yueshengys committed 1 year ago

3b03e549

Update XLA dependency to use revision

a-googler committed 1 year ago

20646eb0

Newer Older