mirror of
https://github.com/ROCm/jax.git
synced 2025-04-18 21:06:06 +00:00
95bc2ba1b9
In the common case (real values) these are all single-expression jaxprs themselves, so putting them out of line just makes things more verbose.
There's no reason to include stuff like this in a jaxpr:
```
cxd:bool[8,16] = pjit[
jaxpr={ lambda ; cxe:f32[8,16]. let
cxf:bool[8,16] = is_finite cxe
in (cxf,) }
name=isfinite
] cxc
```
PiperOrigin-RevId: 587047955