CI: 12/02/24 upstream sync #170
charleshofer
merged
101 commits into
rocm-main
from
ci-daily-sync-02-12-2024
[Pallas] Deprecate dictionary compiler_params in favor of dataclass.
73fa0f48
Add a jaxlib at head build to the cloud-tpu-ci-nightly workflow
34a2f0ca
Adds coverage for spmd-axisname-filtering in shard_map transpose.
c0811c9d
Use with_spec where possible to clean up the code a bit
76356052
[sharding_in_types] Make argmax and argmin work with sharding_in_type…
21f8885a
Add linearization rule for pjit_p
b1d1dcf6
Update XLA dependency to use revision
9f6dbef3
Merge pull request #25070 from jax-ml:pjit-lin-rule
030ee4a1
When host local inputs on all hosts are the same, use `_DeferredShard…
8699f5d9
Fix member access to xla backend. The correct member is `client` inst…
b259fde5
[Mosaic][Easy] - Wire up kernel names to MLIR dump
e53ff2cb
Update XLA dependency to use revision
4d8751bf
Update XLA dependency to use revision
b372ce4b
[pallas:mosaic_gpu] Add test for FragmentedArray.bitcast.
69e3f0d3
Refactor FFI examples to consolidate several examples into one submod…
84a9cba8
[Mosaic GPU] Simplify logic for pointwise splat operands
914600a0
[ROCm] Add rocm version information
e8934b95
Automated Code Change
aa05dc0b
Add abstract mesh context manager to trace_context in the fallback pa…
c35f8b22
[cuda] Bump nvidia-cuda-nvcc-cu12 dependency to 12.6.85
9866372d
[SDY] enable `cpu_shardy` for JAX shard_alike test.
bb1024f3
[SDY] Enable `test_pjit_array_multi_input_multi_output` since Shardy …
066859e6
Update ROCm scripts to match new build.py usage
84dc9bab
Remove _pjit_lower_cached cache. We can simplify the caching of jit a…
deab6fbd
[Mosaic GPU] Support batch dimensions in FA3 MGPU kernel.
107bc96c
drop compute capability check
95029abc
[SDY] remove TODO for enabling Layouts for Shardy post cl/697715276.
f22bafac
Re-factor build CLI to a subcommand based approach
67615126
Merge pull request #25041 from dfm:ffi-example-refactor
788f4935
Add new CI scripts for building JAX artifacts
f7e9f625
Update XLA dependency to use revision
ebea4353
[pallas_mgpu] Allow trees (eg tuples) to be returned from cond_p expr…
ef7df1ae
[mgpu/pallas_mgpu] Pointwise tanh support
c5dc980d
Add sharding argument to reshape since it also takes a `shape` argume…
59e13f81
Create a `null_mesh_context` internal context manager to handle null …
627debc7
[mgpu] Pointwise min
f828f2d7
Merge pull request #25084 from ROCm:ci_rocm_version
024e3314
Use xla_extension_version instead of jaxlib_version
16a5607c
[mosaic_gpu] Fixed unbounded recursion in `FragmentedArray._pointwise`
b6566c80
[AutoPGLE] Explicitly ignore host callback pointers
231967fd
[Pallas TPU] Better error message for lowering `sp.broadcast_to_p`
dc11d402
[AutoPGLE] Fix pgle test after removing pjit cache.
92e18e6d
Update XLA dependency to use revision
e453fa17
Fix a weird interaction with `set_local` and empty tuples passed to it.
6763fcfb
[ROCm] fix jax and wheelhouse relative paths
d30ec2b5
[ROCm] Change run_multi_gpu set opts
694de6b6
[JAX] Add Python binding for building a colocated Python program
bbaec6ea
Add argument to override base docker in dockerfile
8df27664
Update http to https in amd artifactory url.
3d806320
Update XLA dependency to use revision
9c423796
Move `tsl/platform/{build_config,build_config_root,rules_cc}.bzl` to …
10fdee34
Install git before actions/checkout
afcef677
Add a check for return codes of `executor.run` so that we propagate e…
c6866d05
Add new CI script to run Bazel GPU (non-RBE) jobs
13726690
Add a private `set_mesh` API to enter into sharding_in_types mode. Th…
0d2dfea4
Update the render documentation job to use the new self-hosted runners
47d19609
[Mosaic:TPU] Enable broadcast from 1-D vectors
7a2070e7
[mosaic_gpu] Warmup before measuring the running time in `profiler.me…
7f14de04
Integrate LLVM at llvm/llvm-project@b214ca82daee
03b6945e
[mgpu] FragentedArray.foreach() can now optionally return a new array
f3acfa93
[mgpu pallas] Layout iota operation.
8477580d
Fix early exiting when building multiple wheels
d449f12a
[mgpu] Debug print for mlir vectors.
df8ecb97
Merge pull request #25096 from nitins17:update-rocm-ci-scripts
04a4f9bd
Update XLA dependency to use revision
df6758f0
[AutoPGLE] Update fdo_profile comment.
c2c177ee
[Pallas] Fix float -> int casting on Triton backend.
6e72592b
Raise a better error message if anything other than a sequence of int…
cc5036cc
Update XLA dependency to use revision
a212a29d
Add experimental JAX roofline API.
8c521547
Update XLA dependency to use revision
132ad251
Rework custom hermetic python instructions.
b62ca8b1
Merge pull request #25153 from epiqueras:feature/typechecker
bdee4c38
[mgpu] foreach should not try to create an array if it didn't create …
34fe66b0
Reverts cc5036cc18bc585b0d92a4f606956da084effbad
a158e02b
[Mosaic GPU] Add support for fast upcasts of s8 to bf16 for vectors o…
b09b0779
[Mosaic GPU] Avoid double-predication when async_copy predicate is sp…
14ddb819
[mgpu] Added a missed case for debug_print types and raise a proper e…
d5bfafbc
[Pallas][Mosaic GPU] Add support for compressing squeezed dims in asy…
b801539f
[Mosaic GPU] Improve the implementation of max and exp
db158e6c
[Take 2] Raise a better error message if anything other than a sequen…
456dfeb0
Update XLA dependency to use revision
f73de230
Merge pull request #25137 from ROCm:ci_enable_https-upstream
385e2f43
Merge pull request #25091 from gspschmid:gschmid/nvidia-cuda-nvcc-cu1…
6d4278d2
Merge pull request #25130 from ROCm:ci_fix_set_options-upstream
b0df4052
Merge pull request #25128 from ROCm:ci_fix_wheelhouse_relative_paths-…
ab79066b
[Mosaic GPU] Allow contracting ops into FMAs
f10d3eb3
[mgpu] Fixed off-by-one issue in pointwise argument shuffling when le…
ea69401e
[mgpu_pallas] Better support for unsigned integers and floats in iota.
c3c21c74
Add new CI scripts for running Pytests
031c0acf
Update XLA dependency to use revision
47858c4a
Update XLA dependency to use revision
db4b3f29
C++ tree with path API
a1dfdc1d
Update XLA dependency to use revision
e124c051
[Mosaic GPU] Add a bank-conflict checker to tiled transfer + transfer…
bd66f528
Merge pull request #25136 from ROCm:ci_dockerfile_arg_changes-upstream
7b32d882
[jax] Canonicalize dtypes when checking if dtypes present in target d…
5d5b06cf
[Pallas:MGPU] Fix an overly strict precision requirement in tests
aff7714d
Merge branch 'rocm-main' into ci-daily-sync-02-12-2024
58e30455
Update ci-build.yaml to use specific image
97d201e2
Update ci-build.yaml
f8b753cf
JehandadKhan
approved these changes
on 2024-12-02
Reviewers
Assignees
No one assigned
Labels
None yet
Milestone
No milestone
Login to write a write a comment.
Login via GitHub