如何在TensorFlow中用concat/stack/linalg.band_part构造指定分块矩阵？

以下是几种无需手动定义矩阵、也不依赖assign操作（规避图模式限制）的TensorFlow实现方式，均可生成目标设计矩阵：

方法1：基于单位矩阵扩展 #

利用单位矩阵的结构，将每个元素扩展为长度为e的全1/全0向量，再合并维度：

import tensorflow as tf

z = 3
e = 3

# 创建z×z的单位矩阵
identity_mat = tf.eye(z)
# 将单位矩阵的每个元素扩展为e维向量，得到z×z×e的张量
expanded_blocks = tf.expand_dims(identity_mat, axis=-1) * tf.ones((z, z, e), dtype=tf.float32)
# 合并最后两个维度，得到z×(z*e)的目标矩阵
design_matrix = tf.reshape(expanded_blocks, (z, z * e))

print(design_matrix.numpy())

方法2：基于one-hot编码与tile操作 #

通过one-hot编码标记每行的目标块位置，再对每个位置重复e次：

import tensorflow as tf

z = 3
e = 3

# 生成每行对应的块索引：[0, 1, 2]
row_indices = tf.range(z)
# 生成z×z的one-hot矩阵，每行仅对应位置为1
one_hot_mat = tf.one_hot(row_indices, depth=z)
# 对每行的每个元素重复e次，将单个1扩展为连续e个1
design_matrix = tf.tile(one_hot_mat, multiples=[1, e])

print(design_matrix.numpy())

方法3：基于concat拼接块矩阵 #

生成多个z×e的子块矩阵，每个子块仅对应行全为1，再沿列方向拼接：

import tensorflow as tf

z = 3
e = 3

blocks = []
for block_idx in range(z):
    # 创建z×e的子块，仅第block_idx行全为1，其余为0
    block = tf.where(
        tf.equal(tf.range(z)[:, tf.newaxis], block_idx),
        tf.ones((z, e), dtype=tf.float32),
        tf.zeros((z, e), dtype=tf.float32)
    )
    blocks.append(block)
# 沿列拼接所有子块，得到目标矩阵
design_matrix = tf.concat(blocks, axis=1)

print(design_matrix.numpy())

以上三种方法均完全基于TensorFlow的图兼容操作，不会出现assign在图模式中的限制问题，且可轻松扩展到任意z和e的取值。

内容的提问来源于stack exchange，提问作者user1571823