vecs = np.array([[[1, 2, 3, 4],  # Input shape: (2, 3, 4)
                      [5, 6, 7, 8],
                      [9, 8, 7, 6]],
                     [[5, 4, 3, 2],
                      [1, 2, 3, 4],