Some changes to the shift operation to allow it to work optimally regardless of the input and output dimensions.