refactored wavelet_2d_layer