def wrap_reshape_apply_fn(stat_fn, a, b, axis_a, axis_b):
    """
    Reshape two nd-arrays into 2d-arrays, apply function and reshape
    result back.

    Parameters
    ----------
    stat_fn: function
        Function to apply to 2d-arrays
    a: nd-array: nd-array
    b: nd-array
    axis_a: int or list of int
        sample axis
    axis_b: int or list of int
        sample axis

    Returns
    -------
    result: nd-array
        The result reshaped to remaining_dims_a + remaining_dims_b
    """
    if not hasattr(axis_a, "__len__"):
        axis_a = [axis_a]
    if not hasattr(axis_b, "__len__"):
        axis_b = [axis_b]
    other_axis_a = [i for i in range(a.ndim) if i not in axis_a]
    other_axis_b = [i for i in range(b.ndim) if i not in axis_b]
    transposed_topo_a = a.transpose(tuple(other_axis_a) + tuple(axis_a))
    n_stat_axis_a = [a.shape[i] for i in axis_a]
    n_other_axis_a = [a.shape[i] for i in other_axis_a]
    flat_topo_a = transposed_topo_a.reshape(
        np.prod(n_other_axis_a), np.prod(n_stat_axis_a)
    )
    transposed_topo_b = b.transpose(tuple(other_axis_b) + tuple(axis_b))
    n_stat_axis_b = [b.shape[i] for i in axis_b]
    n_other_axis_b = [b.shape[i] for i in other_axis_b]
    flat_topo_b = transposed_topo_b.reshape(
        np.prod(n_other_axis_b), np.prod(n_stat_axis_b)
    )
    assert np.array_equal(n_stat_axis_a, n_stat_axis_b)
    stat_result = stat_fn(flat_topo_a, flat_topo_b)
    topo_result = stat_result.reshape(tuple(n_other_axis_a) + tuple(n_other_axis_b))
    return topo_result

    Reshape two nd-arrays into 2d-arrays, apply function and reshape
    result back.

    Parameters
    ----------
    stat_fn: function
        Function to apply to 2d-arrays
    a: nd-array: nd-array
    b: nd-array
    axis_a: int or list of int
        sample axis
    axis_b: int or list of int
        sample axis

    Returns
    -------
    result: nd-array
        The result reshaped to remaining_dims_a + remaining_dims_b
    
braindecode/util.py