ROCm/jax

Pull Requests Commits

Revert "Merge pull request #297 from ROCm/prep-for-v0.5.1"

charleshofer committed 272 days ago

555d7ef8

Merge pull request #297 from ROCm/prep-for-v0.5.1

Ruturaj4 committed 280 days ago

Verified c07a0e88

Enable upstream CI on release branches (#296)

charleshofer committed 280 days ago

f56eeaf7

Fix auditwheel version issue (#288)

mrodden committed 280 days ago

c38c23c7

Fix numactl-devel name

charleshofer committed 281 days ago

54ce4c89

Fix numa package

charleshofer committed 281 days ago

c9501cb4

Install numa library

charleshofer committed 281 days ago

79e2c71b

Change CODEOWNERS (#237)

charleshofer committed 281 days ago

065bea83

Use bazel for PR tests (#216)

charleshofer committed 281 days ago

e9168054

Prepare for JAX release 0.5.1

yashk2810 committed 303 days ago

07440f4a

Update XLA dependency to use revision

Google-ML-Automation committed 303 days ago

b0cfcb8e

[sharding_in_types] Add more reshape sharding support

yashk2810 committed 303 days ago

7d3c63ed

[mosaic_gpu] Warmup the kernel when doing CUPTI based profiling

superbobry committed 303 days ago

908ff49e

[pallas:mosaic_gpu] Use `{min,max}imumf` instead of `{min,max}numf`

superbobry committed 304 days ago

74b2e020

Update XLA dependency to use revision

Google-ML-Automation committed 304 days ago

e0e5457c

[sharding_in_types] Add sharding rules for the following primitives:

yashk2810 committed 304 days ago

d695aa4c

Update XLA dependency to use revision

Google-ML-Automation committed 305 days ago

a75bea51

[sharding_in_types] Allow `auto_axes` and `explicit_axes` to take numpy arrays, python scalars.

yashk2810 committed 305 days ago

7c4fe2a7

Reverts 6e83de5909b5dad6827c1682726c840c9688b32d

yashk2810 committed 305 days ago

34077851

[sharding_in_types] Make slice and ellipsis work with `.at[...].get(out_sharding=P(...))`

yashk2810 committed 305 days ago

80f18ded

Allow casting to the same axis type

yashk2810 committed 305 days ago

629426f8

Add CODEOWNERS file (#236)

charleshofer committed 305 days ago

Verified 60f51d21

Implement Jax CPU/GPU callbacks with XLA's FFI.

danielsuo committed 305 days ago

6e83de59

roofline: Handle ClosedJaxpr instances.

matthiaskramm committed 305 days ago

b3fcba7c

[JAX] Implement an initial object API for colocated Python

hyeontaek committed 305 days ago

96b7dbab

Reverts 655267609b5589fda8358ae7aaf2eb832036407a

GleasonK committed 305 days ago

6c83d436

Update several remaining lax.linalg primitives to use registration helper functions.

dfm committed 305 days ago

c4418c10

Use new ML Build CUDA images

nitins17 committed 305 days ago

5089fb01

Update Windows docker image to pin it to a sha

nitins17 committed 305 days ago

a0736391

Refactor Jax FFI lowering to prepare for implementing CPU/GPU callbacks using XLA's FFI.

danielsuo committed 306 days ago

2d1bc5c2

Older