scipy 1.15.3__cp312-cp312-win_amd64.whl → 1.16.0rc2__cp312-cp312-win_amd64.whl

scipy/linalg/_matfuncs.py

@@ -8,14 +8,17 @@ import numpy as np
 from numpy import (dot, diag, prod, logical_not, ravel, transpose,
                    conjugate, absolute, amax, sign, isfinite, triu)
 
+from scipy._lib._util import _apply_over_batch
+from scipy._lib.deprecation import _NoValue
+
 # Local imports
-from scipy.linalg import LinAlgError, bandwidth
+from scipy.linalg import LinAlgError, bandwidth, LinAlgWarning
 from ._misc import norm
 from ._basic import solve, inv
 from ._decomp_svd import svd
 from ._decomp_schur import schur, rsf2csf
 from ._expm_frechet import expm_frechet, expm_cond
-from ._matfuncs_sqrtm import sqrtm
+from ._matfuncs_schur_sqrtm import recursive_schur_sqrtm
 from ._matfuncs_expm import pick_pade_structure, pade_UV_calc
 from ._linalg_pythran import _funm_loops  # type: ignore[import-not-found]
 
@@ -95,6 +98,7 @@ def _maybe_real(A, B, tol=None):
 # Matrix functions.
 
 
+@_apply_over_batch(('A', 2))
 def fractional_matrix_power(A, t):
     """
     Compute the fractional power of a matrix.
@@ -141,7 +145,8 @@ def fractional_matrix_power(A, t):
     return scipy.linalg._matfuncs_inv_ssq._fractional_matrix_power(A, t)
 
 
-def logm(A, disp=True):
+@_apply_over_batch(('A', 2))
+def logm(A, disp=_NoValue):
     """
     Compute matrix logarithm.
 
@@ -156,6 +161,11 @@ def logm(A, disp=True):
         Emit warning if error in the result is estimated large
         instead of returning estimated error. (Default: True)
 
+        .. deprecated:: 1.16.0
+            The `disp` argument is deprecated and will be
+            removed in SciPy 1.18.0. The previously returned error estimate
+            can be computed as ``norm(expm(logm(A)) - A, 1) / norm(A, 1)``.
+
     Returns
     -------
     logm : (N, N) ndarray
@@ -196,6 +206,12 @@ def logm(A, disp=True):
            [ 1.,  4.]])
 
     """
+    if disp is _NoValue:
+        disp = True
+    else:
+        warnings.warn("The `disp` argument is deprecated "
+                      "and will be removed in SciPy 1.18.0.",
+                      DeprecationWarning, stacklevel=2)
     A = np.asarray(A)  # squareness checked in `_logm`
     # Avoid circular import ... this is OK, right?
     import scipy.linalg._matfuncs_inv_ssq
@@ -391,6 +407,153 @@ def _exp_sinch(x):
     return lexp_diff
 
 
+def sqrtm(A, disp=_NoValue, blocksize=_NoValue):
+    """
+    Compute, if exists, the matrix square root.
+
+    The matrix square root of ``A`` is a matrix ``X`` such that ``X @ X = A``.
+    Every square matrix is not guaranteed to have a matrix square root, for
+    example, the array ``[[0, 1], [0, 0]]`` does not have a square root.
+
+    Moreover, not every real matrix has a real square root. Hence, for
+    real-valued matrices the return type can be complex if, numerically, there
+    is an eigenvalue on the negative real axis.
+
+    Parameters
+    ----------
+    A : ndarray
+        Input with last two dimensions are square ``(..., n, n)``.
+    disp : bool, optional
+        Print warning if error in the result is estimated large
+        instead of returning estimated error. (Default: True)
+
+        .. deprecated:: 1.16.0
+            The `disp` argument is deprecated and will be
+            removed in SciPy 1.18.0. The previously returned error estimate
+            can be computed as ``norm(X @ X - A, 'fro')**2 / norm(A, 'fro')``
+
+    blocksize : integer, optional
+
+        .. deprecated:: 1.16.0
+            The `blocksize` argument is deprecated as it is unused by the algorithm
+            and will be removed in SciPy 1.18.0.
+
+    Returns
+    -------
+    sqrtm : ndarray
+        Computed matrix squareroot of `A` with same size ``(..., n, n)``.
+
+    errest : float
+        Frobenius norm of the estimated error, ||err||_F / ||A||_F. Only
+        returned, if ``disp`` is set to ``False``. This return argument will be
+        removed in version 1.20.0 and only the sqrtm result will be returned.
+
+        .. deprecated:: 1.16.0
+
+    Notes
+    -----
+    This function uses the Schur decomposition method to compute the matrix
+    square root following [1]_ and for real matrices [2]_. Moreover, note
+    that, there exist matrices that have square roots that are not polynomials
+    in ``A``. For a classical example from [2]_, the matrix satisfies::
+
+            [ a, a**2 + 1]**2     [-1,  0]
+            [-1,       -a]     =  [ 0, -1]
+
+    for any scalar ``a`` but it is not a polynomial in ``-I``. Thus, they will
+    not be found by this function.
+
+    References
+    ----------
+    .. [1] Edvin Deadman, Nicholas J. Higham, Rui Ralha (2013)
+           "Blocked Schur Algorithms for Computing the Matrix Square Root,
+           Lecture Notes in Computer Science, 7782. pp. 171-182.
+           :doi:`10.1016/0024-3795(87)90118-2`
+    .. [2] Nicholas J. Higham (1987) "Computing real square roots of a real
+           matrix", Linear Algebra and its Applications, 88/89:405-430.
+           :doi:`10.1016/0024-3795(87)90118-2`
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from scipy.linalg import sqrtm
+    >>> a = np.array([[1.0, 3.0], [1.0, 4.0]])
+    >>> r = sqrtm(a)
+    >>> r
+    array([[ 0.75592895,  1.13389342],
+           [ 0.37796447,  1.88982237]])
+    >>> r.dot(r)
+    array([[ 1.,  3.],
+           [ 1.,  4.]])
+
+    """
+    if disp is _NoValue:
+        disp = True
+    else:
+        warnings.warn("The `disp` argument is deprecated and will be removed in SciPy "
+                      "1.18.0.",
+                      DeprecationWarning, stacklevel=2)
+    if blocksize is not _NoValue:
+        warnings.warn("The `blocksize` argument is deprecated and will be removed in "
+                      "SciPy 1.18.0.",
+                      DeprecationWarning, stacklevel=2)
+
+    a = np.asarray(A)
+    if a.size == 1 and a.ndim < 2:
+        return np.array([[np.exp(a.item())]])
+
+    if a.ndim < 2:
+        raise LinAlgError('The input array must be at least two-dimensional')
+    if a.shape[-1] != a.shape[-2]:
+        raise LinAlgError('Last 2 dimensions of the array must be square')
+
+    # Empty array
+    if min(*a.shape) == 0:
+        dtype = sqrtm(np.eye(2, dtype=a.dtype)).dtype
+        return np.empty_like(a, dtype=dtype)
+
+    # Scalar case
+    if a.shape[-2:] == (1, 1):
+        return np.emath.sqrt(a)
+
+    if not np.issubdtype(a.dtype, np.inexact):
+        a = a.astype(np.float64)
+    elif a.dtype == np.float16:
+        a = a.astype(np.float32)
+    elif a.dtype.char in 'G':
+        a = a.astype(np.complex128)
+    elif a.dtype.char in 'g':
+        a = a.astype(np.float64)
+
+    if a.dtype.char not in 'fdFD':
+        raise TypeError("scipy.linalg.sqrtm is not supported for the data type"
+                        f" {a.dtype}")
+
+    res, isIllconditioned, isSingular, info = recursive_schur_sqrtm(a)
+    if info < 0:
+        raise LinAlgError(f"Internal error in scipy.linalg.sqrtm: {info}")
+
+    if isSingular or isIllconditioned:
+        if isSingular:
+            msg = ("Matrix is singular. The result might be inaccurate or the"
+                   " array might not have a square root.")
+        else:
+            msg = ("Matrix is ill-conditioned. The result might be inaccurate"
+                   " or the array might not have a square root.")
+        warnings.warn(msg, LinAlgWarning, stacklevel=2)
+
+    if disp is False:
+        try:
+            arg2 = norm(res @ res - A, 'fro')**2 / norm(A, 'fro')
+        except ValueError:
+            # NaNs in matrix
+            arg2 = np.inf
+        return res, arg2
+    else:
+        return res
+
+
+@_apply_over_batch(('A', 2))
 def cosm(A):
     """
     Compute the matrix cosine.
@@ -431,6 +594,7 @@ def cosm(A):
         return expm(1j*A).real
 
 
+@_apply_over_batch(('A', 2))
 def sinm(A):
     """
     Compute the matrix sine.
@@ -471,6 +635,7 @@ def sinm(A):
         return expm(1j*A).imag
 
 
+@_apply_over_batch(('A', 2))
 def tanm(A):
     """
     Compute the matrix tangent.
@@ -510,6 +675,7 @@ def tanm(A):
     return _maybe_real(A, solve(cosm(A), sinm(A)))
 
 
+@_apply_over_batch(('A', 2))
 def coshm(A):
     """
     Compute the hyperbolic matrix cosine.
@@ -549,6 +715,7 @@ def coshm(A):
     return _maybe_real(A, 0.5 * (expm(A) + expm(-A)))
 
 
+@_apply_over_batch(('A', 2))
 def sinhm(A):
     """
     Compute the hyperbolic matrix sine.
@@ -588,6 +755,7 @@ def sinhm(A):
     return _maybe_real(A, 0.5 * (expm(A) - expm(-A)))
 
 
+@_apply_over_batch(('A', 2))
 def tanhm(A):
     """
     Compute the hyperbolic matrix tangent.
@@ -627,6 +795,7 @@ def tanhm(A):
     return _maybe_real(A, solve(coshm(A), sinhm(A)))
 
 
+@_apply_over_batch(('A', 2))
 def funm(A, func, disp=True):
     """
     Evaluate a matrix function specified by a callable.
@@ -721,7 +890,8 @@ def funm(A, func, disp=True):
         return F, err
 
 
-def signm(A, disp=True):
+@_apply_over_batch(('A', 2))
+def signm(A, disp=_NoValue):
     """
     Matrix sign function.
 
@@ -735,6 +905,11 @@ def signm(A, disp=True):
         Print warning if error in the result is estimated large
         instead of returning estimated error. (Default: True)
 
+        .. deprecated:: 1.16.0
+            The `disp` argument is deprecated and will be
+            removed in SciPy 1.18.0. The previously returned error estimate
+            can be computed as ``norm(signm @ signm - signm, 1)``.
+
     Returns
     -------
     signm : (N, N) ndarray
@@ -754,6 +929,13 @@ def signm(A, disp=True):
     array([-1.+0.j,  1.+0.j,  1.+0.j])
 
     """
+    if disp is _NoValue:
+        disp = True
+    else:
+        warnings.warn("The `disp` argument is deprecated "
+                      "and will be removed in SciPy 1.18.0.",
+                      DeprecationWarning, stacklevel=2)
+
     A = _asarray_square(A)
 
     def rounded_sign(x):
@@ -801,6 +983,7 @@ def signm(A, disp=True):
         return S0, errest
 
 
+@_apply_over_batch(('a', 2), ('b', 2))
 def khatri_rao(a, b):
     r"""
     Khatri-rao product

scipy/linalg/_matfuncs_sqrtm.py

@@ -4,24 +4,16 @@ Matrix square root for general matrices and for upper triangular matrices.
 This module exists to avoid cyclic imports.
 
 """
-__all__ = ['sqrtm']
+__all__ = []
 
 import numpy as np
 
-from scipy._lib._util import _asarray_validated
-
 # Local imports
-from ._misc import norm
 from .lapack import ztrsyl, dtrsyl
-from ._decomp_schur import schur, rsf2csf
-from ._basic import _ensure_dtype_cdsz
-
-
 
 class SqrtmError(np.linalg.LinAlgError):
     pass
 
-
 from ._matfuncs_sqrtm_triu import within_block_loop  # noqa: E402
 
 
@@ -113,93 +105,3 @@ def _sqrtm_triu(T, blocksize=64):
 
     # Return the matrix square root.
     return R
-
-
-def sqrtm(A, disp=True, blocksize=64):
-    """
-    Matrix square root.
-
-    Parameters
-    ----------
-    A : (N, N) array_like
-        Matrix whose square root to evaluate
-    disp : bool, optional
-        Print warning if error in the result is estimated large
-        instead of returning estimated error. (Default: True)
-    blocksize : integer, optional
-        If the blocksize is not degenerate with respect to the
-        size of the input array, then use a blocked algorithm. (Default: 64)
-
-    Returns
-    -------
-    sqrtm : (N, N) ndarray
-        Value of the sqrt function at `A`. The dtype is float or complex.
-        The precision (data size) is determined based on the precision of
-        input `A`.
-
-    errest : float
-        (if disp == False)
-
-        Frobenius norm of the estimated error, ||err||_F / ||A||_F
-
-    References
-    ----------
-    .. [1] Edvin Deadman, Nicholas J. Higham, Rui Ralha (2013)
-           "Blocked Schur Algorithms for Computing the Matrix Square Root,
-           Lecture Notes in Computer Science, 7782. pp. 171-182.
-
-    Examples
-    --------
-    >>> import numpy as np
-    >>> from scipy.linalg import sqrtm
-    >>> a = np.array([[1.0, 3.0], [1.0, 4.0]])
-    >>> r = sqrtm(a)
-    >>> r
-    array([[ 0.75592895,  1.13389342],
-           [ 0.37796447,  1.88982237]])
-    >>> r.dot(r)
-    array([[ 1.,  3.],
-           [ 1.,  4.]])
-
-    """
-    A = _asarray_validated(A, check_finite=True, as_inexact=True)
-    if len(A.shape) != 2:
-        raise ValueError("Non-matrix input to matrix function.")
-    if blocksize < 1:
-        raise ValueError("The blocksize should be at least 1.")
-    A, = _ensure_dtype_cdsz(A)
-    keep_it_real = np.isrealobj(A)
-    if keep_it_real:
-        T, Z = schur(A)
-        d0 = np.diagonal(T)
-        d1 = np.diagonal(T, -1)
-        eps = np.finfo(T.dtype).eps
-        needs_conversion = abs(d1) > eps * (abs(d0[1:]) + abs(d0[:-1]))
-        if needs_conversion.any():
-            T, Z = rsf2csf(T, Z)
-    else:
-        T, Z = schur(A, output='complex')
-    failflag = False
-    try:
-        R = _sqrtm_triu(T, blocksize=blocksize)
-        ZH = np.conjugate(Z).T
-        X = Z.dot(R).dot(ZH)
-        dtype = np.result_type(A.dtype, 1j if np.iscomplexobj(X) else 1)
-        X = X.astype(dtype, copy=False)
-    except SqrtmError:
-        failflag = True
-        X = np.empty_like(A)
-        X.fill(np.nan)
-
-    if disp:
-        if failflag:
-            print("Failed to find a square root.")
-        return X
-    else:
-        try:
-            arg2 = norm(X.dot(X) - A, 'fro')**2 / norm(A, 'fro')
-        except ValueError:
-            # NaNs in matrix
-            arg2 = np.inf
-
-        return X, arg2

scipy/linalg/_procrustes.py

@@ -3,12 +3,14 @@ Solve the orthogonal Procrustes problem.
 
 """
 import numpy as np
+from scipy._lib._util import _apply_over_batch
 from ._decomp_svd import svd
 
 
 __all__ = ['orthogonal_procrustes']
 
 
+@_apply_over_batch(('A', 2), ('B', 2))
 def orthogonal_procrustes(A, B, check_finite=True):
     """
     Compute the matrix solution of the orthogonal (or unitary) Procrustes problem.

scipy/linalg/_sketches.py

@@ -6,8 +6,8 @@
 import numpy as np
 
 from scipy._lib._util import (check_random_state, rng_integers,
-                              _transition_to_rng)
-from scipy.sparse import csc_matrix
+                              _transition_to_rng, _apply_over_batch)
+from scipy.sparse import csc_matrix, issparse
 
 __all__ = ['clarkson_woodruff_transform']
 
@@ -68,8 +68,8 @@ def clarkson_woodruff_transform(input_matrix, sketch_size, rng=None):
 
     Parameters
     ----------
-    input_matrix : array_like
-        Input matrix, of shape ``(n, d)``.
+    input_matrix : array_like, shape (..., n, d)
+        Input matrix.
     sketch_size : int
         Number of rows for the sketch.
     rng : `numpy.random.Generator`, optional
@@ -174,5 +174,16 @@ def clarkson_woodruff_transform(input_matrix, sketch_size, rng=None):
     166.58473879945151
 
     """
-    S = cwt_matrix(sketch_size, input_matrix.shape[0], rng=rng)
-    return S.dot(input_matrix)
+    if issparse(input_matrix) and input_matrix.ndim > 2:
+        message = "Batch support for sparse arrays is not available."
+        raise NotImplementedError(message)
+
+    S = cwt_matrix(sketch_size, input_matrix.shape[-2], rng=rng)
+    # Despite argument order (required by decorator), this is  S @ input_matrix
+    # Can avoid _batch_dot when gh-22153 is resolved.
+    return S @ input_matrix if input_matrix.ndim <= 2 else _batch_dot(input_matrix, S)
+
+
+@_apply_over_batch(('input_matrix', 2))
+def _batch_dot(input_matrix, S):
+    return S @ input_matrix

scipy/linalg/_solvers.py

@@ -12,6 +12,7 @@ import warnings
 import numpy as np
 from numpy.linalg import inv, LinAlgError, norm, cond, svd
 
+from scipy._lib._util import _apply_over_batch
 from ._basic import solve, solve_triangular, matrix_balance
 from .lapack import get_lapack_funcs
 from ._decomp_schur import schur
@@ -27,6 +28,7 @@ __all__ = ['solve_sylvester',
            'solve_continuous_are', 'solve_discrete_are']
 
 
+@_apply_over_batch(('a', 2), ('b', 2), ('q', 2))
 def solve_sylvester(a, b, q):
     """
     Computes a solution (X) to the Sylvester equation :math:`AX + XB = Q`.
@@ -106,12 +108,12 @@ def solve_sylvester(a, b, q):
     y = scale*y
 
     if info < 0:
-        raise LinAlgError("Illegal value encountered in "
-                          "the %d term" % (-info,))
+        raise LinAlgError(f"Illegal value encountered in the {-info} term")
 
     return np.dot(np.dot(u, y), v.conj().transpose())
 
 
+@_apply_over_batch(('a', 2), ('q', 2))
 def solve_continuous_lyapunov(a, q):
     """
     Solves the continuous Lyapunov equation :math:`AX + XA^H = Q`.
@@ -245,6 +247,7 @@ def _solve_discrete_lyapunov_bilinear(a, q):
     return solve_lyapunov(b.conj().transpose(), -c)
 
 
+@_apply_over_batch(('a', 2), ('q', 2))
 def solve_discrete_lyapunov(a, q, method=None):
     """
     Solves the discrete Lyapunov equation :math:`AXA^H - X + Q = 0`.
@@ -336,6 +339,7 @@ def solve_discrete_lyapunov(a, q, method=None):
     return x
 
 
+@_apply_over_batch(('a', 2), ('b', 2), ('q', 2), ('r', 2), ('e', 2), ('s', 2))
 def solve_continuous_are(a, b, q, r, e=None, s=None, balanced=True):
     r"""
     Solves the continuous-time algebraic Riccati equation (CARE).
@@ -541,6 +545,7 @@ def solve_continuous_are(a, b, q, r, e=None, s=None, balanced=True):
     return (x + x.conj().T)/2
 
 
+@_apply_over_batch(('a', 2), ('b', 2), ('q', 2), ('r', 2), ('e', 2), ('s', 2))
 def solve_discrete_are(a, b, q, r, e=None, s=None, balanced=True):
     r"""
     Solves the discrete-time algebraic Riccati equation (DARE).

scipy/linalg/_special_matrices.py

@@ -3,6 +3,7 @@ import warnings
 
 import numpy as np
 from numpy.lib.stride_tricks import as_strided
+from scipy._lib._util import _apply_over_batch
 
 
 __all__ = ['toeplitz', 'circulant', 'hankel',
@@ -269,6 +270,7 @@ def hadamard(n, dtype=int):
     return H
 
 
+@_apply_over_batch(("f", 1), ("s", 1))
 def leslie(f, s):
     """
     Create a Leslie matrix.
@@ -279,24 +281,22 @@ def leslie(f, s):
 
     Parameters
     ----------
-    f : (..., N,) array_like
+    f : (N,) array_like
         The "fecundity" coefficients.
-    s : (..., N-1,) array_like
-        The "survival" coefficients. The length of each slice of `s` (along the last
-        axis) must be one less than the length of `f`, and it must be at least 1.
+    s : (N-1,) array_like
+        The "survival" coefficients. The length of `s` must be one less
+        than the length of `f`, and it must be at least 1.
 
     Returns
     -------
-    L : (..., N, N) ndarray
+    L : (N, N) ndarray
         The array is zero except for the first row,
         which is `f`, and the first sub-diagonal, which is `s`.
-        For 1-D input, the data-type of the array will be the data-type of
+        The data-type of the array will be the data-type of
         ``f[0]+s[0]``.
 
     Notes
     -----
-    .. versionadded:: 0.8.0
-
     The Leslie matrix is used to model discrete-time, age-structured
     population growth [1]_ [2]_. In a population with `n` age classes, two sets
     of parameters define a Leslie matrix: the `n` "fecundity coefficients",
@@ -304,11 +304,6 @@ def leslie(f, s):
     class, and the `n` - 1 "survival coefficients", which give the
     per-capita survival rate of each age class.
 
-    N-dimensional input are treated as a batches of coefficient arrays: each
-    slice along the last axis of the input arrays is a 1-D coefficient array,
-    and each slice along the last two dimensions of the output is the
-    corresponding Leslie matrix.
-
     References
     ----------
     .. [1] P. H. Leslie, On the use of matrices in certain population
@@ -337,12 +332,6 @@ def leslie(f, s):
         raise ValueError("The length of s must be at least 1.")
 
     n = f.shape[-1]
-
-    if f.ndim > 1 or s.ndim > 1:
-        from scipy.stats._resampling import _vectorize_statistic
-        _leslie_nd = _vectorize_statistic(leslie)
-        return np.moveaxis(_leslie_nd(f, s, axis=-1), [0, 1], [-2, -1])
-
     tmp = f[0] + s[0]
     a = np.zeros((n, n), dtype=tmp.dtype)
     a[0] = f
@@ -406,9 +395,9 @@ def kron(a, b):
 
 def block_diag(*arrs):
     """
-    Create a block diagonal matrix from provided arrays.
+    Create a block diagonal array from provided arrays.
 
-    Given the inputs `A`, `B` and `C`, the output will have these
+    For example, given 2-D inputs `A`, `B` and `C`, the output will have these
     arrays arranged on the diagonal::
 
         [[A, 0, 0],
@@ -417,15 +406,17 @@ def block_diag(*arrs):
 
     Parameters
     ----------
-    A, B, C, ... : array_like, up to 2-D
-        Input arrays.  A 1-D array or array_like sequence of length `n` is
-        treated as a 2-D array with shape ``(1,n)``.
+    A, B, C, ... : array_like
+        Input arrays.  A 1-D array or array_like sequence of length ``n`` is
+        treated as a 2-D array with shape ``(1, n)``. Any dimensions before
+        the last two are treated as batch dimensions; see :ref:`linalg_batch`.
 
     Returns
     -------
     D : ndarray
-        Array with `A`, `B`, `C`, ... on the diagonal. `D` has the
-        same dtype as `A`.
+        Array with `A`, `B`, `C`, ... on the diagonal of the last two
+        dimensions. `D` has the same dtype as the result type of the
+        inputs.
 
     Notes
     -----
@@ -433,7 +424,8 @@ def block_diag(*arrs):
     block diagonal matrix.
 
     Empty sequences (i.e., array-likes of zero size) will not be ignored.
-    Noteworthy, both [] and [[]] are treated as matrices with shape ``(1,0)``.
+    Noteworthy, both ``[]`` and ``[[]]`` are treated as matrices with shape
+    ``(1,0)``.
 
     Examples
     --------
@@ -470,18 +462,16 @@ def block_diag(*arrs):
         arrs = ([],)
     arrs = [np.atleast_2d(a) for a in arrs]
 
-    bad_args = [k for k in range(len(arrs)) if arrs[k].ndim > 2]
-    if bad_args:
-        raise ValueError("arguments in the following positions "
-                         f"have dimension greater than 2: {bad_args}")
-
-    shapes = np.array([a.shape for a in arrs])
+    batch_shapes = [a.shape[:-2] for a in arrs]
+    batch_shape = np.broadcast_shapes(*batch_shapes)
+    arrs = [np.broadcast_to(a, batch_shape + a.shape[-2:]) for a in arrs]
     out_dtype = np.result_type(*[arr.dtype for arr in arrs])
-    out = np.zeros(np.sum(shapes, axis=0), dtype=out_dtype)
+    block_shapes = np.array([a.shape[-2:] for a in arrs])
+    out = np.zeros(batch_shape + tuple(np.sum(block_shapes, axis=0)), dtype=out_dtype)
 
     r, c = 0, 0
-    for i, (rr, cc) in enumerate(shapes):
-        out[r:r + rr, c:c + cc] = arrs[i]
+    for i, (rr, cc) in enumerate(block_shapes):
+        out[..., r:r + rr, c:c + cc] = arrs[i]
         r += rr
         c += cc
     return out