Fix Bug in HaarWaveletTransform3D, Added correct output packing with k=8

[yasinzaii]

The final einops.rearrange in HaarWaveletTransform3D.forward incorrectly used k=c, which is only true when C==8. Even though it seems like it works I am guessing because we have just specified incorrect k, so it compensates the rest with c. The first dimension of the stacked outputs is 8*(B*C), so the proper factorization is (b k c) with k=8.

[yasinzaii]

The order was also incorrect - The order for the outputs is (b c k) before rearranging as we can clearly see that we are appending in the k dimension which is the fastest changing dimension and should come at last. but then we rearrange it into (k,c) meaning the fastest changing dimension in now c. and for each wavelet band such as LL we stack all of the Channels.

When we look at the inverse wavelet transform this is exactly what we do. we divide the input into 8 chunks. Which actually means that each chunk will have all of those channels. so again here channel is the fastest changing dimension and now the inverse transform fully complements the forward transform.

[qqingzheng]

Thank you so much!!! I will check it soon!