Fixes to infer_shape for tensordot

Author: lshaw8317

src/blosc2/shape_utils.py

@@ -542,6 +542,8 @@ def _lookup_value(self, node):
             return self.shapes.get(node.id, None)
         elif isinstance(node, ast.Constant):
             return node.value
+        elif isinstance(node, ast.Tuple):
+            return tuple(e.value for e in node.elts)
         else:
             return None
 

tests/ndarray/test_lazyexpr.py

@@ -1714,10 +1714,14 @@ def test_lazylinalg():
     np.testing.assert_array_almost_equal(out[()], npres)
 
     # --- tensordot ---
-    out = blosc2.lazyexpr("tensordot(A, B, axes=1)")
+    out = blosc2.lazyexpr("tensordot(A, B, axes=1)")  # test with int axes
     npres = np.tensordot(npA, npB, axes=1)
     assert out.shape == npres.shape
     np.testing.assert_array_almost_equal(out[()], npres)
+    out = blosc2.lazyexpr("tensordot(A, B, axes=((1,) , (0,)))")  # test with tuple axes
+    npres = np.tensordot(npA, npB, axes=((1,), (0,)))
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
 
     # --- vecdot ---
     out = blosc2.lazyexpr("vecdot(x, y)")