How to solve batch_size problem?

Question

Python

def construct_autoregressive_mask(target):
    """
        Args: Original target of word ids, shape [batch, seq_len]
        Returns: a mask of shape [batch, seq_len, seq_len].
    """
    batch_size, seq_len = target.shape.as_list()
    
    print(batch_size)

    tri_matrix = np.zeros((seq_len, seq_len))
    tri_matrix[np.tril_indices(seq_len)] = 1
    
    mask = tf.convert_to_tensor(tri_matrix, dtype=tf.float32)

    masks = tf.tile(tf.expand_dims(mask, 0), (batch_size, 1, 1))  # copies
    return masks

Throws an error:
TypeError: Failed to convert object of type to Tensor. Contents: (None, 1, 1). Consider casting elements to a supported type.
But if you initialize batch_size (For example, immediately indicate that it is equal to 64 or create a placeholder), there is no error.
Is it possible to solve the problem without these actions?

Answer 1

Use tf.shape(target) to get a tensor with the shape of another tensor whose size can change from step to step.
Tensorflow, by the way, can automatically bring forms. If we add tensor A with shape (K, L, M) and tensor B with shape (1, L, M), then TF will automatically copy B along the first axis K times, and when adding tensors with shapes (K, 1) and ( 1, L) will also copy both and the result will be of the form (K, L): https://www.tensorflow.org/xla/broadcasting