Revert change to next_power_of_two.

Reverts 168f30a

PiperOrigin-RevId: 615722408

Author: chr1sj0nes

jax/_src/pallas/utils.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Pallas utility functions."""
+import math
 from jax import lax
 from jax._src import core as jax_core
 from jax._src.util import split_list
@@ -44,7 +45,9 @@ def strides_from_shape(shape: tuple[int, ...]) -> tuple[int, ...]:
 
 
 def next_power_of_2(x: int) -> int:
-  return 2**x.bit_length()
+  if x == 0:
+    return 1
+  return int(2 ** math.ceil(math.log2(x)))
 
 
 def pattern_match_scan_to_fori_loop(