Custom Orthonormal Kernel Constraint

[GPUser]

Hello, I'm trying to implement a custom kernel constraint function that will convert the weight kernel to an orthonormal basis, such that the output neurons of that layer would be linearly independent of each other. However, my implementation of the function seems to be really inefficient, to the point where I don't think it would be practical (or even possible) to train with the constraint. Here is the code:

# Design an orthonormal weight constraint that uses the Gram-Schmidt process
# to keep all output neural activities linearly independent (and well-scaled).
constraint_orthonormal_2d <- function(w) {
  # Dimensions of kernel: (rows, cols, in_chan, out_chan).
  w <- k_eval(w)
  N_out <- dim(w)[4]
  for (i in 1:N_out) {
    # Normalize the weight vector for output neuron i.
    w[,,,i] <- k_eval(k_l2_normalize(w[,,,i]))
    
    # Remove the projection onto weight vector i from all other weight vectors.
    if (i == N_out) break
    for (j in (i+1):N_out) {
      w[,,,j] <- k_eval(w[,,,j] - w[,,,i] * k_sum(w[,,,i] * w[,,,j]))
    }
  }
  return(k_variable(w))
}

A big part of the problem seems to be the subsetting operations, although I feel that the for-loops are also inappropriate for this type of function. The k_eval() and k_variable() functions were my workaround attempt for the fact that you can't subset Tensors. RStudio Keras has brief documentation on creating custom kernel constraints, but it doesn't look like there is a way to perform a constraint function that requires subsetting the weight tensor.

Unfortunately, the algorithm for creating an orthonormal basis (the Gram-Schmidt process) requires iteration. Is there another way to implement an orthonormal constraint in Keras? Could an alternate formulation using recursive function calls work?