Fix Array import in segment_utils

PiperOrigin-RevId: 611536788
google · Feb 29, 2024 · 3537d8f · 3537d8f
1 parent 073798e
commit 3537d8f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/rax/_src/segment_utils.py b/rax/_src/segment_utils.py
@@ -19,7 +19,9 @@
 import jax
 import jax.numpy as jnp
 
-from rax._src.types import Array
+from rax._src import types
+
+Array = types.Array
 
 
 def same_segment_mask(segments: Array) -> Array: