Tensorflow dataset element shuffle within specified range

Question

How do I shuffle the elements of tf.data.Dataset within a certain range. Having an input array, with shape = (10,), in the first 5 elements would be shuffled within the first 5 places and the next 5 withing the next 5 places, and so on.

ex: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] ==> [3, 4, 2, 5, 1, 7, 10, 9, 6, 8]

Below is the function I implemented, which works on numpy arrays, but not on tensorflow data tensors.

Could someone please help me out with advice?

def indexShifting(my_array):
    lim = my_array.shape[0]
    shift = 5 # limits the range of shuffle
    rng = np.random.default_rng()

    for i in range(floor(lim/shift)): # 10/5
        rng.shuffle(my_array[(i*shift):((i+1)*shift)])
    
    return my_array

a = np.arange(1,11)
> array([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10])

indexShifting(a)
> array([ 3,  4,  2,  5,  1,  7, 10,  9,  6,  8])

With tf.data.Dataset, there are a few errors I receive.

1. len is not well defined for symbolic Tensors.
2. 'tensorflow.python.framework.ops.EagerTensor' object does not support item assignment

This is dummy data.

a = np.random.permutation(10)
b = np.random.permutation(10)
c = np.random.permutation(10)

X = np.vstack((a,b,c))
y = np.expand_dims(np.ones(3), axis=1)

train_ds = tf.data.Dataset.from_tensor_slices((X, y))

train_ds = train_ds.map(lambda x, y: (indexShifting(x), y))

TF = 2.5.0 Python = 3.8.10

Any help and advice would help me out a lot. Thanks :)