如何在TensorFlow/numpy中实现层的权重与输入节点置换索引视图？

问题背景 #

给定以下置换索引与张量：

import numpy as np

pa = [2,0,4,3,1,5] # 节点置换索引
pw = [0,3,4,1,2,5] # 权重置换索引
a = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
w = np.array([0.7, 0.6, 0.5, 0.4, 0.3, 0.2])

需要实现类似ap = a[pa]的语义，但要求置换后的结果是原张量的引用而非副本——修改置换后的张量元素时，原张量对应位置的值同步更新，且要在TensorFlow中实现该功能，用于任意权重共享场景：不同目标节点以不同顺序引用同一组权重，同时保证反向传播时所有权重实例都能被正确更新。

核心结论 #

Numpy确实不支持非连续索引的视图（因为置换后元素内存不连续，无法直接映射），但TensorFlow基于计算图的符号化机制，可以实现等效的引用效果，完全满足权重共享与反向传播需求。

实现方法 #

方法1：用tf.gather+tf.scatter_update模拟引用

tf.gather会创建基于原张量的符号化引用，反向传播时梯度自动回流到原变量；修改置换后的"视图"时，通过tf.scatter_update直接操作原变量，置换结果会自动同步。

示例代码：

import tensorflow as tf

# 定义可训练的原始权重变量
w = tf.Variable([0.7, 0.6, 0.5, 0.4, 0.3, 0.2])
pw = [0,3,4,1,2,5] # 权重置换索引

# 获取置换后的符号化引用（非副本）
wp = tf.gather(w, pw)

# 模拟修改wp[0]，同步更新原w
# wp[0]对应原w的索引是pw[0]，直接修改原变量对应位置
update_op = tf.scatter_update(w, pw[0], 0.123)

# 验证效果
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print("修改前w:", sess.run(w))
    print("修改前wp:", sess.run(wp))
    sess.run(update_op)
    print("修改后w:", sess.run(w))
    print("修改后wp:", sess.run(wp))
    # 验证：w[pw[0]]与wp[0]值一致

批量修改置换后元素时，可传入索引列表批量更新：

# 修改wp的第0、2位，对应原w的pw[0]、pw[2]位置
indices_in_w = [pw[0], pw[2]]
new_values = [0.123, 0.456]
batch_update_op = tf.scatter_update(w, indices_in_w, new_values)

方法2：自定义置换层实现权重共享

如果用于神经网络层，可以自定义层维护共享权重，前向传播返回置换后的权重，反向传播自动更新原变量。

示例代码：

class PermutedWeightLayer(tf.keras.layers.Layer):
    def __init__(self, permutation, **kwargs):
        super().__init__(**kwargs)
        self.permutation = permutation
    
    def build(self, input_shape):
        # 初始化权重（若要共享，直接赋值外部变量即可）
        self.w = self.add_weight(shape=(input_shape[-1],),
                                initializer='glorot_uniform',
                                trainable=True)
        super().build(input_shape)
    
    def call(self, inputs):
        # 前向传播返回置换后的权重
        permuted_w = tf.gather(self.w, self.permutation)
        return tf.multiply(inputs, permuted_w)

# 共享权重示例
shared_w = tf.Variable(tf.random.normal((6,)))

# 两个不同置换的层，共用同一组权重
layer1 = PermutedWeightLayer([2,0,4,3,1,5])
layer1.w = shared_w

layer2 = PermutedWeightLayer([0,3,4,1,2,5])
layer2.w = shared_w

反向传播时，两个层的梯度会统一更新shared_w，实现权重共享。

关键说明 #

TensorFlow中没有Numpy式的内存视图，但通过符号化计算图的引用机制，tf.gather生成的张量始终与原变量绑定，反向传播梯度自动回流；修改操作直接作用于原变量即可同步所有置换后的引用，完全满足权重共享的需求。

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