ROCm/jax

Pull Requests Commits

Update socket transfers to call new transfer library APIs.

emilyfertig committed 186 days ago

877bee73

Update ml_dtypes to 0.5.1 to align with JAX and TensorFlow

apivovarov committed 186 days ago

21b52fb5

Export `jax.P` as an API which is an alias for `jax.sharding.PartitionSpec`.

yashk2810 committed 186 days ago

0e852e28

[JAX] Fallback to deprecated field output_memory_colors for now, due to a bug in libtpu.

subhankarshah committed 186 days ago

52890321

If shape is empty and out_sharding is None, just use the sharding of key in `random.uniform`. Fixes https://github.com/jax-ml/jax/issues/29694

yashk2810 committed 186 days ago

51b354f9

[mosaic] Removed `kernel` and `backend` parameters from `tpu_custom_call` APIs

superbobry committed 186 days ago

518eb970

remove use of nanobind for _export pass with shardy

liepieshov committed 186 days ago

5f953a9d

[pallas] Fixed the type of `Mesh.backend`

superbobry committed 186 days ago

be329aa4

Add JAX documentation for psend and precv API

Google-ML-Automation committed 186 days ago

c94ab8f4

Merge pull request #29507 from jakevdp:tracer-repr

Google-ML-Automation committed 186 days ago

8516d28b

Have tpu7x support in Jax.

Google-ML-Automation committed 186 days ago

2b050070

Add a hash/equality to JaxprEqnContext.

hawkinsp committed 186 days ago

406fb6d5

Move `jax._src.cudnn.*` sources to their own build rule

Jake VanderPlas committed 186 days ago

ecb0ddef

Rollback "Custom batching rule for Cholesky that ensures symmetrize."

btaba committed 186 days ago

3dacc957

Add very basic support for shard_map + unreduced.

yashk2810 committed 186 days ago

69796d34

[JAX] Disable psend precv shard_map tests temporarily.

Google-ML-Automation committed 186 days ago

78e57d01

Merge pull request #29616 from jakevdp:dlpack-supported-dtypes

Google-ML-Automation committed 186 days ago

7c06e18b

Use more concise tracer reprs by default.

jakevdp committed 186 days ago

90609d47

Merge pull request #29696 from jakevdp:delete-interpolate

Google-ML-Automation committed 186 days ago

38e6888f

Allow empty mesh in lhs/rhs in dot_general. Fixes https://github.com/jax-ml/jax/issues/29683

yashk2810 committed 186 days ago

d585311b

Removed a few old jaxlib version guards

superbobry committed 186 days ago

70b39706

Deprecate jax.dtypes.SUPPORTED_DTYPES and add is_supported_dtype()

jakevdp committed 186 days ago

86eb71ae

Remove type declarations for deprecated APIs

Jake VanderPlas committed 186 days ago

502a6d89

Fix corner case check in 2x2x2 VF mesh creation.

bshahr committed 186 days ago

209e2393

[jax:pjit] Add jaxlib version guards to `ArrayPjitTest.jit_mul_sum_sharding_preserved`.

danielsuo committed 186 days ago

47cb97e0

Custom batching rule for Cholesky that ensures symmetrize.

Google-ML-Automation committed 186 days ago

ebba464b

Remove obsolete submodule jax._src.scipy.interpolate

jakevdp committed 186 days ago

a6294a06

[Pallas][NFC] Add documentation for some CompilerParams.

bythew3i committed 186 days ago

cd93c2d7

[Mosaic GPU] Use mgpu.Barrier instead of mgpu.TMABarrier for the MMA barrier

apaszke committed 186 days ago

6daf870a

[Mosaic GPU] Use the 64-bit mov instruction while moving 64-bit timer values

apaszke committed 186 days ago

28cfc5e6

Older