`jax.experimental.host_callback` has been deprecated since March 2024
(JAX version 0.4.26). Now we set the default value of the `--jax_host_callback_legacy` configuration value to `True`, which means that if your code uses `jax.experimental.host_callback` APIs, those API calls will be implemented in terms of the new `jax.experimental.io_callback` API.
If this breaks your code, for a very limited time, you can set the `--jax_host_callback_legacy` to `True`. Soon we will remove that configuration option, so you should instead transition to using the new JAX callback APIs.
See https://github.com/google/jax/issues/20385 for a discussion.
PiperOrigin-RevId: 681004255
This is a second attempt at this change. The first one was rolled back because of reported failures.
Reverts 411928b9668570bbc3795522aba94cece6894881
PiperOrigin-RevId: 680943744
Historically, tests that only ran on GPUs were placed in `OpsExtraTest`, while general tests were in `OpsTest`. However, this separation may cause us to miss issues that should be addressed on TPUs as well. Going forward, all tests will be unified in `OpsTest`, and any tests that fail on TPUs will be skipped individually using `skipTest`. This will help us better track and address TPU-specific failures.
PiperOrigin-RevId: 680747902
We still have this temporary check in apply vector layout, but in infer vector layout, instead of throwing error, we should just reset offset to zero. Because some ops which has relaxed this restriction might be passed as input for un-relaxed ops and cause failure.
PiperOrigin-RevId: 680706301
This fixes the failure in elementwise rule of apply vector layout pass.
If the condition scalar is static, it will be simplified to corresponding vector from true value and false value by MLIR.
If the condition scalar is dynamic, we want to use vselect over scf.if anyway. Because latter creates a inner region.
PiperOrigin-RevId: 680674560
This changes makes it so that the refs users receive inside their kernels have shapes
matching their block specs. However, the refs are not actually plain refs, but transformed
references that begin with the fully transformed abstract ref and then stack the inverse
of the transformation stack on top of it. This means that all primitives that take in refs
can also see the sequence of transforms the user applied in the block spec, which lets us
verify e.g. that the inputs to WGMMA are correctly tiled, even though their user-visible
shape remains 2D. We should be able to use the same trick in the future to propagate tiling
and better infer the layouts for loads and stores.
PiperOrigin-RevId: 680520185
* Delete custom_object_test, since it is disabled and has been ever since jax.Array was introduced in JAX 0.4.0.
* custom_linear_solve_test was over-sharded, leading to some shards not having any test cases. Even unsharded it completes in under 65s on every platform we have.
* config_test and pallas splash attention mask test only tested helpers and didn't need a TPU.
PiperOrigin-RevId: 679711664
