import numpy as np


def multiprod(A, B):
    """Vectorized matrix-matrix multiplication.

    A and B are
    assumed to be arrays containing M matrices, that is, A and B have
    dimensions A: (M, N, P), B:(M, P, Q). multiprod multiplies each matrix
    in A with the corresponding matrix in B, using matrix multiplication.
    so multiprod(A, B) has dimensions (M, N, Q).
    """

    # First check if we have been given just one matrix
    if len(np.shape(A)) == 2:
        return np.dot(A, B)

    # Approx 5x faster, only supported by numpy version >= 1.6:
    return np.einsum('ijk,ikl->ijl', A, B)


def multitransp(A):
    """Vectorized matrix transpose.

    A is assumed to be an array containing M matrices, each of which has
    dimension N x P.
    That is, A is an M x N x P array. Multitransp then returns an array
    containing the M matrix transposes of the matrices in A, each of which will
    be P x N.
    """
    # First check if we have been given just one matrix
    if A.ndim == 2:
        return A.T
    return np.transpose(A, (0, 2, 1))


def multihconj(A):
    """Vectorized matrix Hermitian conjugate."""
    return np.conjugate(multitransp(A))


def multisym(A):
    """Vectorized matrix symmetrization.

    Given an array ``A`` of matrices (represented as an array of shape ``(k, n,
    n)``), returns a version of ``A`` with each matrix symmetrized, i.e.,
    every matrix ``A[i]`` satisfies ``A[i] == A[i].T``.
    """
    return 0.5 * (A + multitransp(A))


def multiskew(A):
    """Vectorized matrix skew-symmetrization.

    Same as :func:`multisym`, but returns an array where each matrix ``A[i]``
    is skew-symmetric, i.e., the components of ``A`` satisfy ``A[i] ==
    -A[i].T``.
    """
    return 0.5 * (A - multitransp(A))


def multieye(k, n):
    """Array of ``k`` ``n x n`` identity matrices."""
    return np.tile(np.eye(n), (k, 1, 1))


def multilog(A, pos_def=False):
    """Vectorized matrix logarithm."""
    if not pos_def:
        raise NotImplementedError

    w, v = np.linalg.eigh(A)
    w = np.expand_dims(np.log(w), axis=-1)
    return multiprod(v, w * multitransp(v))


def multiexp(A, sym=False):
    """Vectorized matrix exponential."""
    if not sym:
        raise NotImplementedError

    w, v = np.linalg.eigh(A)
    w = np.expand_dims(np.exp(w), axis=-1)
    return multiprod(v, w * multitransp(v))

@@ -29,20 +29,58 @@
 tiprod(A
-, B)
+: np.ndarray, B: np.ndarray) -%3E np.ndarray
 :%0A    %22%22
@@ -131,22 +131,309 @@
    A
- and B are%0A   
+rgs:%0A        A: The first matrix.%0A        B: The second matrix.%0A%0A    Returns:%0A        The matrix (or more precisely array of matrices) corresponding to%0A        the matrix product vectorized over the first dimension of %60%60A%60%60 and%0A        %60%60B%60%60 (if %60%60A.ndim == 2%60%60).%0A%0A    The matrices %60%60A%60%60 and %60%60B%60%60 are
  ass
@@ -453,33 +453,41 @@
 rays containing 
-M
+%60%60k%60%60%0A   
  matrices, that 
@@ -485,81 +485,108 @@
 es, 
-that is, A and B have%0A    dimensions A: (M, N, P), B:(M, P, Q). multiprod
+i.e., %60%60A%60%60 and %60%60B%60%60 have shape %60%60(k, m, n)%60%60 and %60%60(k, n, p)%60%60,%0A    respectively.%0A    The function
  mul
@@ -608,17 +608,17 @@
 trix
-%0A    in A
+ in %60%60A%60%60
  wit
@@ -647,91 +647,96 @@
 trix
- in B, using matrix multiplication.%0A    so multiprod(A, B) has dimensions (M, N, Q)
+%0A    in %60%60B%60%60 along the first dimension.%0A    The resulting array has shape %60%60(k, m, p)%60%60
 .%0A  
@@ -745,187 +745,55 @@
 %22%22%22%0A
-%0A
     
-# First check if we have been given just one matrix%0A    if len(np.shape(A)) == 2:%0A        return np.dot(A, B)%0A%0A    # Approx 5x faster, only supported by numpy version %3E= 1.6:
+if A.ndim == 2:%0A        return np.dot(A, B)
 %0A   
@@ -810,17 +810,17 @@
 .einsum(
-'
+%22
 ijk,ikl-
@@ -827,9 +827,9 @@
 %3Eijl
-'
+%22
 , A,
@@ -1163,64 +1163,8 @@
 %22%22%22%0A
-    # First check if we have been given just one matrix%0A
     

Add docstring for 'multiprod'