Fix numpy 2 errors

PiperOrigin-RevId: 665367904
google-research · Aug 21, 2024 · cde981d · cde981d
1 parent 726f13e
commit cde981d
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 8 deletions.
diff --git a/dataclass_array/array_dataclass_test.py b/dataclass_array/array_dataclass_test.py
@@ -520,13 +520,13 @@ def test_absolute_axis(xnp: enp.NpModule):
   assert p._to_absolute_axis((0, 1, -1)) == (0, 1, 3)
   assert p._to_absolute_axis((-1, -2)) == (3, 2)
 
-  with pytest.raises(np.AxisError):
+  with pytest.raises(np.exceptions.AxisError):
     assert p._to_absolute_axis(4)
 
-  with pytest.raises(np.AxisError):
+  with pytest.raises(np.exceptions.AxisError):
     assert p._to_absolute_axis(-5)
 
-  with pytest.raises(np.AxisError):
+  with pytest.raises(np.exceptions.AxisError):
     assert p._to_absolute_axis((0, 4))
 
 

diff --git a/dataclass_array/shape_parsing.py b/dataclass_array/shape_parsing.py
@@ -75,7 +75,7 @@ def singleton(cls) -> ShapeParser:
 
   def __init__(self):
     grammar_path = file_utils.dca_path() / 'shape_grammar.lark'
-    self.parser = lark.Lark(grammar_path.read_text())
+    self.parser = lark.Lark(grammar_path.read_text(), ambiguity='resolve')
     self.transformer = _TreeShapeTransformer()
 
   def parse(self, shape_str: str) -> _ShapeAst:

diff --git a/dataclass_array/utils/np_utils.py b/dataclass_array/utils/np_utils.py
@@ -157,7 +157,7 @@ def to_absolute_axis(axis: Axes, *, ndim: int) -> Axes:
     return tuple(range(ndim))
   elif isinstance(axis, int):
     if axis >= ndim or axis < -ndim:
-      raise enp.lazy.np.AxisError(
+      raise enp.lazy.np.exceptions.AxisError(
           axis=axis,
           ndim=ndim,
           # msg_prefix=

diff --git a/dataclass_array/utils/np_utils_test.py b/dataclass_array/utils/np_utils_test.py
@@ -55,13 +55,13 @@ def test_to_absolute_axis():
   assert np_utils.to_absolute_axis((0, 1, -1), ndim=4) == (0, 1, 3)
   assert np_utils.to_absolute_axis((-1, -2), ndim=4) == (3, 2)
 
-  with pytest.raises(np.AxisError):
+  with pytest.raises(np.exceptions.AxisError):
     assert np_utils.to_absolute_axis(4, ndim=4)
 
-  with pytest.raises(np.AxisError):
+  with pytest.raises(np.exceptions.AxisError):
     assert np_utils.to_absolute_axis(-5, ndim=4)
 
-  with pytest.raises(np.AxisError):
+  with pytest.raises(np.exceptions.AxisError):
     assert np_utils.to_absolute_axis((0, 4), ndim=4)