tensorflow - 将输入与矩阵相乘后，批量大小已更改

Question

我正在 TensorFlow 中编写一个自定义层来计算 PCA。当我将形状（None, 150, 300）的输入提供给该层时，我会收到形状为（150, 150, 100）且n_components = 100的输出。

我怎样才能保持批量大小=无？我想要的输出是：（无，150，100）。

def call(self, input_data):
            covariance_matrix = tfp.stats.covariance(input_data)
        
            self.eigen_values, self.eigen_vectors = tf.linalg.eig(covariance_matrix)
            self.eigen_values = tf.cast(self.eigen_values, dtype=tf.float32)
            self.eigen_vectors = tf.cast(self.eigen_vectors, dtype=tf.float32)
    
            self.projection_matrix = tf.transpose(tf.transpose(self.eigen_vectors)[:][-self.n_components:])
        
            output_data = tf.matmul(input_data, self.projection_matrix)
            
            return output_data