Support updated typing in numpy 2.0

Author: ntjohnson1

pyttb/cp_als.py

@@ -185,7 +185,7 @@ def cp_als(  # noqa: PLR0912,PLR0913,PLR0915
 
     # Reduce dimorder to only those modes we will optimize
     dimorder_in = dimorder  # save for output
-    dimorder = [d for d in dimorder if d in optdims]
+    dimorder = [int(d) for d in dimorder if d in optdims]
 
     # Store the last MTTKRP result to accelerate fitness computation
     U_mttkrp = np.zeros((input_tensor.shape[dimorder[-1]], rank))

pyttb/cp_apr.py

@@ -917,8 +917,8 @@ def tt_cp_apr_pqnr(  # noqa: PLR0912,PLR0913,PLR0915
                 delg = np.zeros((rank, lbfgsMem))
                 rho = np.zeros((lbfgsMem,))
                 lbfgsPos = 0
-                m_rowOLD = np.empty(())
-                gradOLD = np.empty(())
+                m_rowOLD = np.empty((), dtype=m_row.dtype)
+                gradOLD = np.empty((), dtype=m_row.dtype)
 
                 # Iteratively solve the row subproblem with projected quasi-Newton steps
                 for i in range(maxinneriters):
@@ -983,8 +983,8 @@ def tt_cp_apr_pqnr(  # noqa: PLR0912,PLR0913,PLR0915
                     isRowNOTconverged[jj] = 1
 
                     # Update the L-BFGS approximation.
-                    tmp_delm = m_row - m_rowOLD
-                    tmp_delg = gradM - gradOLD
+                    tmp_delm: np.ndarray = m_row - m_rowOLD
+                    tmp_delg: np.ndarray = gradM - gradOLD
                     tmp_delm_dot = tmp_delm.dot(tmp_delg.transpose())
                     if not np.any(tmp_delm_dot == 0):
                         tmp_rho = 1 / tmp_delm_dot

pyttb/gcp/fg_est.py

@@ -162,7 +162,7 @@ def estimate_helper(
     ndim = subs.shape[1]
 
     # Create exploded U's from the model factor matrices
-    Uexp = [np.empty(())] * ndim
+    Uexp = [np.empty((), dtype=factors[0].dtype)] * ndim
     for k in range(ndim):
         Uexp[k] = factors[k][subs[:, k], :]
 

pyttb/gcp/samplers.py

@@ -272,7 +272,7 @@ def nonzeros(
 
     # Select nonzeros
     if samples == nnz:
-        nidx = np.arange(0, nnz)
+        nidx: np.ndarray = np.arange(0, nnz, dtype=int)
     elif with_replacement or samples < nnz:
         nidx = np.random.choice(nnz, size=samples, replace=with_replacement)
     else:

pyttb/hosvd.py

@@ -129,7 +129,7 @@ def hosvd(  # noqa: PLR0912,PLR0913,PLR0915
 
         # Shrink!
         if sequential:
-            Y = Y.ttm(factor_matrices[k].transpose(), k)
+            Y = Y.ttm(factor_matrices[k].transpose(), int(k))
     # Extract final core
     if sequential:
         G = Y

pyttb/ktensor.py

@@ -1192,7 +1192,9 @@ def mask(self, W: Union[ttb.tensor, ttb.sptensor]) -> np.ndarray:
             vals = vals + tmpvals
         return vals
 
-    def mttkrp(self, U: Union[ktensor, Sequence[np.ndarray]], n: int) -> np.ndarray:
+    def mttkrp(
+        self, U: Union[ktensor, Sequence[np.ndarray]], n: Union[int, np.integer]
+    ) -> np.ndarray:
         """
         Matricized tensor times Khatri-Rao product for :class:`pyttb.ktensor`.
 

pyttb/pyttb_utils.py

@@ -816,7 +816,7 @@ def get_index_variant(indices: IndexType) -> IndexVariant:
 
 
 def get_mttkrp_factors(
-    U: Union[ttb.ktensor, Sequence[np.ndarray]], n: int, ndims: int
+    U: Union[ttb.ktensor, Sequence[np.ndarray]], n: Union[int, np.integer], ndims: int
 ) -> Sequence[np.ndarray]:
     """Apply standard checks and type conversions for mttkrp factors"""
     if isinstance(U, ttb.ktensor):

pyttb/sptensor.py

@@ -415,7 +415,7 @@ def allsubs(self) -> np.ndarray:
 
         # Generate each column of the subscripts in turn
         for n in range(0, self.ndims):
-            i = o.copy()
+            i: list[np.ndarray] = o.copy()
             i[n] = np.expand_dims(np.arange(0, self.shape[n]), axis=1)
             s[:, n] = np.squeeze(ttb.khatrirao(*i))
 
@@ -807,15 +807,19 @@ def to_sptenmat(
         csize = np.array(self.shape)[cdims]
 
         if rsize.size == 0:
-            ridx = np.zeros((self.nnz, 1))
+            ridx: np.ndarray[tuple[int, ...], np.dtype[Any]] = np.zeros(
+                (self.nnz, 1), dtype=int
+            )
         elif self.subs.size == 0:
             ridx = np.array([], dtype=int)
         else:
             ridx = tt_sub2ind(rsize, self.subs[:, rdims])
         ridx = ridx.reshape((ridx.size, 1)).astype(int)
 
         if csize.size == 0:
-            cidx = np.zeros((self.nnz, 1))
+            cidx: np.ndarray[tuple[int, ...], np.dtype[Any]] = np.zeros(
+                (self.nnz, 1), dtype=int
+            )
         elif self.subs.size == 0:
             cidx = np.array([], dtype=int)
         else:
@@ -1243,7 +1247,9 @@ def mask(self, W: sptensor) -> np.ndarray:
         vals[matching_indices] = self.vals[matching_indices]
         return vals
 
-    def mttkrp(self, U: Union[ttb.ktensor, Sequence[np.ndarray]], n: int) -> np.ndarray:
+    def mttkrp(
+        self, U: Union[ttb.ktensor, Sequence[np.ndarray]], n: Union[int, np.integer]
+    ) -> np.ndarray:
         """
         Matricized tensor times Khatri-Rao product using the
         :class:`pyttb.sptensor`. This is an efficient form of the matrix
@@ -1312,7 +1318,7 @@ def mttkrp(self, U: Union[ttb.ktensor, Sequence[np.ndarray]], n: int) -> np.ndar
                 else:
                     Z.append(np.array([]))
             # Perform ttv multiplication
-            ttv = self.ttv(Z, exclude_dims=n)
+            ttv = self.ttv(Z, exclude_dims=int(n))
             # TODO is is possible to hit the float condition here?
             if isinstance(ttv, float):  # pragma: no cover
                 V[:, r] = ttv

pyttb/sumtensor.py

@@ -318,7 +318,9 @@ def innerprod(
             result += part.innerprod(other)
         return result
 
-    def mttkrp(self, U: Union[ttb.ktensor, List[np.ndarray]], n: int) -> np.ndarray:
+    def mttkrp(
+        self, U: Union[ttb.ktensor, List[np.ndarray]], n: Union[int, np.integer]
+    ) -> np.ndarray:
         """
         Matricized tensor times Khatri-Rao product. The matrices used in the
         Khatri-Rao product are passed as a :class:`pyttb.ktensor` (where the

pyttb/tensor.py

@@ -8,7 +8,6 @@
 
 import logging
 from collections.abc import Iterable
-from functools import partial
 from itertools import combinations_with_replacement, permutations
 from math import factorial, prod
 from typing import (
@@ -20,6 +19,7 @@
     Sequence,
     Tuple,
     Union,
+    cast,
     overload,
 )
 
@@ -734,9 +734,8 @@ def issymmetric(  # noqa: PLR0912
         sz = np.array(self.shape)
 
         if grps is None:
-            grps = np.arange(0, n)
-
-        if len(grps.shape) == 1:
+            grps = np.arange(0, n)[None, :]
+        elif len(grps.shape) == 1:
             grps = np.array([grps])
 
         # Substantially different routines are called depending on whether the user
@@ -903,7 +902,9 @@ def mask(self, W: tensor) -> np.ndarray:
         # Extract those non-zero values
         return self.data[tuple(wsubs.transpose())]
 
-    def mttkrp(self, U: Union[ttb.ktensor, Sequence[np.ndarray]], n: int) -> np.ndarray:
+    def mttkrp(
+        self, U: Union[ttb.ktensor, Sequence[np.ndarray]], n: Union[int, np.integer]
+    ) -> np.ndarray:
         """
         Matricized tensor times Khatri-Rao product. The matrices used in the
         Khatri-Rao product are passed as a :class:`pyttb.ktensor` (where the
@@ -1272,7 +1273,9 @@ def squeeze(self) -> Union[tensor, float]:
         else:
             idx = np.where(shapeArray > 1)
             if idx[0].size == 0:
-                return self.data.item()
+                # Why is item annotated as str?
+                single_item: float = cast(float, self.data.item())
+                return single_item
             return ttb.tensor(np.squeeze(self.data))
 
     def symmetrize(  # noqa: PLR0912,PLR0915
@@ -1518,7 +1521,7 @@ def ttm(
         newshape = np.array(
             [p, *list(shape[range(0, n)]), *list(shape[range(n + 1, self.ndims)])]
         )
-        Y_data = np.reshape(newdata, newshape, order=self.order)
+        Y_data: np.ndarray = np.reshape(newdata, newshape, order=self.order)
         Y_data = np.transpose(Y_data, np.argsort(order))
         return ttb.tensor(Y_data, copy=False)
 
@@ -1774,7 +1777,7 @@ def ttsv(
 
             # extract scalar if needed
             if len(y) == 1:
-                y = y.item()
+                return cast(float, y.item())
 
             return y
         assert False, "Invalid value for version; should be None, 1, or 2"
@@ -2600,7 +2603,10 @@ def tenones(shape: Shape, order: Union[Literal["F"], Literal["C"]] = "F") -> ten
      [1. 1. 1.]
      [1. 1. 1.]]
     """
-    ones = partial(np.ones, order=order)
+
+    def ones(shape: Tuple[int, ...]) -> np.ndarray:
+        return np.ones(shape, order=order)
+
     return tensor.from_function(ones, shape)
 
 
@@ -2634,7 +2640,10 @@ def tenzeros(shape: Shape, order: Union[Literal["F"], Literal["C"]] = "F") -> te
      [0. 0. 0.]
      [0. 0. 0.]]
     """
-    zeros = partial(np.zeros, order=order)
+
+    def zeros(shape: Tuple[int, ...]) -> np.ndarray:
+        return np.zeros(shape, order=order)
+
     return tensor.from_function(zeros, shape)