jax
002078b3
- Only infer sharding from input in full_like (in eager mode) if the input's sharding is concrete i.e. does not contain an AbstractMesh
Go
Login via GitHub
Home
Pricing
FAQ
Install
Login
via GitHub
Commit
View On
GitHub
Commit
238 days ago
Only infer sharding from input in full_like (in eager mode) if the input's sharding is concrete i.e. does not contain an AbstractMesh PiperOrigin-RevId: 766962130
References
#26604 - [pallas:triton] Fix atomic min/max lowering for uint and float types
#29227 - Only infer sharding from input in full_like (in eager mode) if the input's sharding is concrete i.e. does not contain an AbstractMesh
#477 - CI: 06/20/25 upstream sync
#31381 - Remove old ROCm build code
#31768 - [ROCm] Support lowering through PJRT_Triton_Extension
#32115 - Relax version requirements for ROCm Jax Plugin wheels
#33157 - [ROCm] Resolve undefined behavior in bitshift unit test
#579 - Create rocm-test-requirements.txt
#581 - Fix/pallas tests shared memory
#584 - Use plain bazel to test jax, use hermetic rocm dependency
#585 - update a test for checking zero ROCm GPU event
#34135 - [ROCm] update to test if there are GPU events when doing profiling on…
#591 - Adding run_pytest_rocm.sh
#592 - Upstream ci download jax rocm wheels
#34304 - [ROCm] Add ROCm platform support for psend/precv collective operations
#599 - Use xla from rocm-jaxlib-v0.8.0 release branch
#34450 - [ROCm] Upstream ROCm CI Nightly Wheel Testing
#34470 - [ROCm] Enabled testTridiagonal for ROCm devices.
#34474 - [ROCm] Added support for GESVDJ on ROCm devices
#34494 - [ROCm] Enabled ROCm devices to default to Jacobi SVD on smaller matrices.
#632 - Add script for running bazel tests on ROCm
#638 - Support wheels as build deps
#34561 - [ROCm] Enable ToeplitzSymmetricConstruction and condition number tests
#34567 - [ROCm]Unskip testMultivariateNormalSingularCovariance on ROCm
#34568 - [ROCm] Skip test_tridiagonal_solve on ROCm
#34575 - [ROCm] Enabled the Triton Pallas tests to run for ROCm.
#34577 - [ROCm] Enabled RNN unit test "test_no_workspace_overflow" for ROCm.
#34578 - [ROCm] Unskip supported dtypes for testConvolutionsPreferredElementType on ROCm
#34598 - Implement approx_tanh for ROCm using OCML tanh function
#34600 - Enable scaled_matmul tests on ROCm
#34601 - Enable neural network tests on ROCm
#34602 - Fix and enable Pallas ops tests on ROCm
#34603 - Enable miscellaneous tests on ROCm
#34615 - [ROCm] Skip test_batch_axis_sharding_jvp on ROCm
#34641 - [ROCm] Add support to dynamically set test deps for jax as an external repo
#34675 - [ROCm] Update Skip Reason Outputs
#664 - update with ROCm profiling tests with kernel_details
#670 - Test pytest abort plugin
#671 - Rocm jaxlib v0.9.0
#34689 - [ROCm] Fix ROCm GPU architecture detection and route to Triton backend
Author
yashk2810
Committer
Google-ML-Automation
Parents
8519fd2e
Loading