Avoid floating point error in jnp.log2 with x64 enabled

[obackhouse]

Use of jnp.log2 can cause floating point errors on x64:

>>> import jax
>>> jax.config.update("jax_enable_x64", True)
>>> import jax.numpy as jnp
>>> arr = jnp.zeros((2,) * 6, dtype=jnp.float64)
>>> jnp.log2(arr.size)
Array(6., dtype=float64)
>>> int(jnp.log2(arr.size))
5