ROCm/jax

Pull Requests Commits

Update http to https in amd artifactory url.

Ruturaj4 committed 1 year ago

3d806320

Fix a weird interaction with `set_local` and empty tuples passed to it.

yashk2810 committed 1 year ago

6763fcfb

Update XLA dependency to use revision

belitskiy committed 1 year ago

e453fa17

[AutoPGLE] Fix pgle test after removing pjit cache.

Google-ML-Automation committed 1 year ago

92e18e6d

[Pallas TPU] Better error message for lowering `sp.broadcast_to_p`

ayaka14732 committed 1 year ago

dc11d402

[AutoPGLE] Explicitly ignore host callback pointers

Google-ML-Automation committed 1 year ago

231967fd

[mosaic_gpu] Fixed unbounded recursion in `FragmentedArray._pointwise`

superbobry committed 1 year ago

b6566c80

Use xla_extension_version instead of jaxlib_version

Google-ML-Automation committed 1 year ago

16a5607c

Merge pull request #25084 from ROCm:ci_rocm_version

Google-ML-Automation committed 1 year ago

024e3314

[mgpu] Pointwise min

cperivol committed 1 year ago

f828f2d7

Create a `null_mesh_context` internal context manager to handle null contexts properly.

yashk2810 committed 1 year ago

627debc7

Add sharding argument to reshape since it also takes a `shape` argument for the output shape

yashk2810 committed 1 year ago

59e13f81

[mgpu/pallas_mgpu] Pointwise tanh support

cperivol committed 1 year ago

c5dc980d

[pallas_mgpu] Allow trees (eg tuples) to be returned from cond_p expressions.

cperivol committed 1 year ago

ef7df1ae

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

ebea4353

Add new CI scripts for building JAX artifacts

nitins17 committed 1 year ago

f7e9f625

Merge pull request #25041 from dfm:ffi-example-refactor

Google-ML-Automation committed 1 year ago

788f4935

Re-factor build CLI to a subcommand based approach

nitins17 committed 1 year ago

67615126

[SDY] remove TODO for enabling Layouts for Shardy post cl/697715276.

Varcho committed 1 year ago

f22bafac

drop compute capability check

Google-ML-Automation committed 1 year ago

95029abc

[Mosaic GPU] Support batch dimensions in FA3 MGPU kernel.

justinjfu committed 1 year ago

107bc96c

Remove _pjit_lower_cached cache. We can simplify the caching of jit as we have downstream caches and a cpp cache too.

yashk2810 committed 1 year ago

deab6fbd

[SDY] Enable `test_pjit_array_multi_input_multi_output` since Shardy conflict resolution is now complete.

Varcho committed 1 year ago

066859e6

[SDY] enable `cpu_shardy` for JAX shard_alike test.

Varcho committed 1 year ago

bb1024f3

Add abstract mesh context manager to trace_context in the fallback path too (which will be deleted after jax 0.4.36 release)

yashk2810 committed 1 year ago

c35f8b22

Automated Code Change

Google-ML-Automation committed 1 year ago

aa05dc0b

[ROCm] Add rocm version information

Ruturaj4 committed 1 year ago

e8934b95

[Mosaic GPU] Simplify logic for pointwise splat operands

apaszke committed 1 year ago

914600a0

Refactor FFI examples to consolidate several examples into one submodule.

dfm committed 1 year ago

84a9cba8

[pallas:mosaic_gpu] Add test for FragmentedArray.bitcast.

petebu committed 1 year ago

69e3f0d3

Older