flax

1fc7bfe9 - Fix temp sample implicit rank promo