I am trying to develop a non-trainable custom PCA layer, where, starting from a float input of size equal to 256, the output is a two-dimensional vector obtained from the PCA application inside the network as last level. Prior to this layer I need to have one Dense layer. Below is my PCA implementation via tensorflow and the custom level:

class TF_PCA:
    def __init__(self, X):
        self.X = X
        self.u = None
        self.singular_values = None
        self.sigma = None

    def fit(self):
        # Perform SVD
        self.singular_values, self.u, _ = tf.linalg.svd(self.X)

        # Create sigma matrix
        self.sigma = tf.linalg.diag(self.singular_values)

    def reduce(self, n_dimensions=None, keep_info=None):
        if keep_info:
            # Normalize singular values
            normalized_singular_values = self.singular_values / sum(self.singular_values)

            # Create the aggregated ladder of kept information per dimension
            ladder = np.cumsum(normalized_singular_values)

            # Get the first index which is above the given information threshold
            index = next(idx for idx, value in enumerate(ladder) if value >= keep_info) + 1
            n_dimensions = index

        # Cut out the relevant part from sigma
        sigma = tf.slice(self.sigma, [0, 0], [self.X.shape[1], n_dimensions])

        # PCA
        pca = tf.matmul(self.u, sigma)
        return pca 

class CustomPCALayer(tf.keras.layers.Layer):
  def __init__(self, num_outputs):
    super(CustomPCALayer, self).__init__()
    self.num_outputs = num_outputs
    self.total = tf.Variable(initial_value=tf.zeros((num_outputs,)), trainable=False)

  def call(self, inputs):
    tf_pca = TF_PCA(inputs)
    tf_pca.fit()
    pca = tf_pca.reduce(n_dimensions=self.num_outputs)
    return tf.convert_to_tensor(pca, dtype=tf.float32)

Once I implemented this I also implemented the neural network:

in_dim = (256,)
out_dim = 2

def build_network(inputs, out_dim):
  x = Dense(256)(inputs)
  x = Dropout(0.2)(x)
  x = CustomPCALayer(out_dim)(x)  # non-trainable layer
  return x

inputs = Input(shape=in_dim)
network = build_network(inputs, out_dim)
model = Model(inputs=inputs, outputs=[network], name="network")

opt = keras.optimizers.Adam(learning_rate=0.0001)

model.compile(optimizer=optimizer, 
              loss="mse", 
              metrics=[tf.keras.metrics.CosineSimilarity()])

As input and output I have something like this:

x = tf.random.uniform((300, 256))
y = tf.random.uniform((300, 2))

If I try to train, it immediately fails with the error NotImplementedError: SVD gradient has not been implemented for input with unknown inner matrix shape.

I would like to understand what is wrong with my architecture in order to solve the error and have a non-trainable layer to perform PCA