flax

6d224f1d - [typing] add assertion to avoid pyrefly error

[typing] add assertion to avoid pyrefly error `sampling_state.logits_buffer` is annotated as `jnp.ndarray | None`, so strict typing requires this assertion. This is not an issue at runtime, because the `None` case would result in an uncaught `AttributeError` in the next line. PiperOrigin-RevId: 960267676