Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Pallas TPU] Better error message for lowering
sp.broadcast_to_p
`sp.broadcast_to_p` is a GPU-specific primitive, but it mistakenly appears in TPU lowerings. This PR improves the error message to reflect this. As an example, currently, users will hit this error when doing: ``` def kernel(x_ref, o_ref): m, n = 32, 8 x = pl.load(x_ref, (jnp.arange(m, dtype=jnp.int32)[:, None], jnp.arange(n, dtype=jnp.int32)[None])) o_ref[...] = x ``` PiperOrigin-RevId: 700290975
- Loading branch information