I have a matrix called U_complex with size (None, 92), where None is equivalent to the batch_size used during the training (let’s say it’s 1000), I want to create another matrix called mask with size (None, 128), then fill the mask (None, 0:4:128) and mask(None, 1:32:128) by zeros while all other values should be filled by the matrix of U_complex.

I used the following code:

N = 128
def create_mask(N, U_complex):
    mask = tf.zeros(N, dtype=tf.complex64) # Initialize mask with zeros

    # Indices to zero out
    indices_zeros = tf.concat([tf.range(0, N, 4), tf.range(1, N, 32)], axis=0)
    updates_zeros = tf.zeros(tf.shape(indices_zeros)[0], dtype=tf.complex64)
    mask = tf.tensor_scatter_nd_update(mask, tf.expand_dims(indices_zeros, 1), updates_zeros)

    # Create a boolean mask for non-zero indices
    all_indices = tf.range(N)
    non_zero_indices = ~tf.reduce_any(tf.equal(tf.expand_dims(all_indices, 1), tf.expand_dims(indices_zeros, 0)), axis=1)

    # Reshape non_zero_indices for broadcasting: (1, N)
    non_zero_indices = tf.reshape(non_zero_indices, [1, N])
    non_zero_indices = tf.broadcast_to(non_zero_indices, [tf.shape(U_complex)[0], N])

    # Update the mask directly using the non-zero indices
    mask = tf.where(non_zero_indices, U_complex, mask)

    return mask

but, I get an error when using mask = tf.where(non_zero_indices, U_complex, mask) saying that non_zero_indices and complex must have the same size ! however they should not be the same. How can I solve that issue?