jax

d2d090b8 - Add `select_n` to `roofline`.

Add `select_n` to `roofline`. This implementation is almost the same as a `unary` op, but it's possible to pass a variable number of `*cases` to `select_n`. PiperOrigin-RevId: 774891852