Reverts 168f30a

PiperOrigin-RevId: 615715466
jax-ml · Mar 14, 2024 · 2fd80e3 · 2fd80e3
1 parent 4a35c12
commit 2fd80e3
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/jax/_src/pallas/utils.py b/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(