def is_matrix(A):
    """Checks if an array is a matrix

    Args:
        A (jax.numpy.ndarray): A JAX array


    Returns:
        bool: True if the array is a matrix, False otherwise.
    """
    return A.ndim == 2