def expm(a: ndarray, method: str = "pade") -> ndarray:
    """
    Matrix exponential.

    Returns exp(A) for each (M x M) slice into a multi-dimensional
    array, assumed to be of shape (..., M, M);

    By default Pade (implicit) implementation is used.
    However, explicit Taylor(deg = 18) implementation can be used,
    by supplying additional flag `use_explicit = True`.

    Parameters
    ----------
    a : (..., M, M) array_like
        Input matrix or multi-dimensional array of shape (..., M, M).

    method : String method selector to use explicit ('taylor')
        or implicit ('pade'); default = 'pade'.

    Returns
    -------
    exp(A): matrix exponential of input, or a matrix exponential
        for each slice in the input.

    Notes
    -----
    Implicit Pade implementation is more stable but more computationally
    intensive than explicit Taylor, which is less stable when matrix norm is
    big enough. Also, Taylor can be slightly more performant for matrices of
    small enough norms, but more memory consuming.

    See Also
    --------
    scipy.linalg.expm

    Availability
    --------
    Multiple GPUs, Multiple CPUs
    """

    if a.ndim < 2 or a.shape[-1] != a.shape[-2] or a.size <= 1:
        raise ValueError(f"Invalid input shape for expm: {a.shape}")

    output = zeros_like(a)

    m_info = get_machine()
    num_PEs = m_info.count()

    # run implicit (Pade) method by default:
    #
    if method == "pade":
        expm_func = expm_impl
    elif method == "taylor":
        expm_func = expm_expl
    else:
        raise ValueError(f"Method {method} not supported.")

    if num_PEs < 2:
        for idx in np.ndindex(a.shape[:-2]):
            mdeg, s = expm_func(a[idx], output[idx])
    else:
        for idx in np.ndindex(a.shape[:-2]):
            flat_index = np.ravel_multi_index(idx, a.shape[:-2])

            # assign work to multiple GPUs in round-robin way:
            #
            findx = int(flat_index)
            with m_info[findx % num_PEs]:
                mdeg, s = expm_func(a[idx], output[idx])

    return output

    Matrix exponential.

    Returns exp(A) for each (M x M) slice into a multi-dimensional
    array, assumed to be of shape (..., M, M);

    By default Pade (implicit) implementation is used.
    However, explicit Taylor(deg = 18) implementation can be used,
    by supplying additional flag `use_explicit = True`.

    Parameters
    ----------
    a : (..., M, M) array_like
        Input matrix or multi-dimensional array of shape (..., M, M).

    method : String method selector to use explicit ('taylor')
        or implicit ('pade'); default = 'pade'.

    Returns
    -------
    exp(A): matrix exponential of input, or a matrix exponential
        for each slice in the input.

    Notes
    -----
    Implicit Pade implementation is more stable but more computationally
    intensive than explicit Taylor, which is less stable when matrix norm is
    big enough. Also, Taylor can be slightly more performant for matrices of
    small enough norms, but more memory consuming.

    See Also
    --------
    scipy.linalg.expm

    Availability
    --------
    Multiple GPUs, Multiple CPUs
    
cupynumeric/linalg/linalg.py