ROCm/jax

Pull Requests Commits

Update socket transfers to call new transfer library APIs.

emilyfertig committed 1 year ago

877bee73

Update ml_dtypes to 0.5.1 to align with JAX and TensorFlow

apivovarov committed 1 year ago

21b52fb5

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

yashk2810 committed 1 year ago

0e852e28

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

subhankarshah committed 1 year 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 1 year ago

51b354f9

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

superbobry committed 1 year ago

518eb970

remove use of nanobind for _export pass with shardy

liepieshov committed 1 year ago

5f953a9d

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

superbobry committed 1 year ago

be329aa4

Add JAX documentation for psend and precv API

Google-ML-Automation committed 1 year ago

c94ab8f4

Merge pull request #29507 from jakevdp:tracer-repr

Google-ML-Automation committed 1 year ago

8516d28b

Have tpu7x support in Jax.

Google-ML-Automation committed 1 year ago

2b050070

Add a hash/equality to JaxprEqnContext.

hawkinsp committed 1 year ago

406fb6d5

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

Jake VanderPlas committed 1 year ago

ecb0ddef

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

btaba committed 1 year ago

3dacc957

Add very basic support for shard_map + unreduced.

yashk2810 committed 1 year ago

69796d34

[JAX] Disable psend precv shard_map tests temporarily.

Google-ML-Automation committed 1 year ago

78e57d01

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

Google-ML-Automation committed 1 year ago

7c06e18b

Use more concise tracer reprs by default.

jakevdp committed 1 year ago

90609d47

Merge pull request #29696 from jakevdp:delete-interpolate

Google-ML-Automation committed 1 year ago

38e6888f

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

yashk2810 committed 1 year ago

d585311b

Removed a few old jaxlib version guards

superbobry committed 1 year ago

70b39706

Deprecate jax.dtypes.SUPPORTED_DTYPES and add is_supported_dtype()

jakevdp committed 1 year ago

86eb71ae

Remove type declarations for deprecated APIs

Jake VanderPlas committed 1 year ago

502a6d89

Fix corner case check in 2x2x2 VF mesh creation.

bshahr committed 1 year ago

209e2393

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

danielsuo committed 1 year ago

47cb97e0

Custom batching rule for Cholesky that ensures symmetrize.

Google-ML-Automation committed 1 year ago

ebba464b

Remove obsolete submodule jax._src.scipy.interpolate

jakevdp committed 1 year ago

a6294a06

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

bythew3i committed 1 year ago

cd93c2d7

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

apaszke committed 1 year ago

6daf870a

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

apaszke committed 1 year ago

28cfc5e6

Older