Improve `axis=` functionality for low-level functions.

[JesseLivezey]

Some of the functions only work for 2d arrays, i.e., (n_time, n_channels). In principle, they should all work for (n_time, ...).