jax

11a45138 - [Mosaic GPU] Explicitly check the smem size so that we get a good error rather than a cryptic cuda error