ROCm/jax

Pull Requests Commits

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

charleshofer committed 341 days ago

555d7ef8

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

Ruturaj4 committed 349 days ago

Verified c07a0e88

Enable upstream CI on release branches (#296)

charleshofer committed 349 days ago

f56eeaf7

Fix auditwheel version issue (#288)

mrodden committed 350 days ago

c38c23c7

Fix numactl-devel name

charleshofer committed 350 days ago

54ce4c89

Fix numa package

charleshofer committed 350 days ago

c9501cb4

Install numa library

charleshofer committed 350 days ago

79e2c71b

Change CODEOWNERS (#237)

charleshofer committed 350 days ago

065bea83

Use bazel for PR tests (#216)

charleshofer committed 350 days ago

e9168054

Prepare for JAX release 0.5.1

yashk2810 committed 1 year ago

07440f4a

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

b0cfcb8e

[sharding_in_types] Add more reshape sharding support

yashk2810 committed 1 year ago

7d3c63ed

[mosaic_gpu] Warmup the kernel when doing CUPTI based profiling

superbobry committed 1 year ago

908ff49e

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

superbobry committed 1 year ago

74b2e020

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

e0e5457c

[sharding_in_types] Add sharding rules for the following primitives:

yashk2810 committed 1 year ago

d695aa4c

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

a75bea51

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

yashk2810 committed 1 year ago

7c4fe2a7

Reverts 6e83de5909b5dad6827c1682726c840c9688b32d

yashk2810 committed 1 year ago

34077851

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

yashk2810 committed 1 year ago

80f18ded

Allow casting to the same axis type

yashk2810 committed 1 year ago

629426f8

Add CODEOWNERS file (#236)

charleshofer committed 1 year ago

Verified 60f51d21

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

danielsuo committed 1 year ago

6e83de59

roofline: Handle ClosedJaxpr instances.

matthiaskramm committed 1 year ago

b3fcba7c

[JAX] Implement an initial object API for colocated Python

hyeontaek committed 1 year ago

96b7dbab

Reverts 655267609b5589fda8358ae7aaf2eb832036407a

GleasonK committed 1 year ago

6c83d436

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

dfm committed 1 year ago

c4418c10

Use new ML Build CUDA images

nitins17 committed 1 year ago

5089fb01

Update Windows docker image to pin it to a sha

nitins17 committed 1 year ago

a0736391

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

danielsuo committed 1 year ago

2d1bc5c2

Older