Ways to implement a unit-norm neural network weight constraint

Question

In the popular neural network framework Keras one can constrain the learned weights in a LinearLayer to have a unit norm. For example:

model.add(Dense(500, input_dim=2, activation='relu', kernel_constraint=unit_norm()))

One can think of various motivations for doing this in practice, so it is not just a hypothetical question. Also note that we are talking about constraining the learned weights, not the output (the latter was discussed in another post.

How would one implement this in Mathematica?

• The NetTrain Method specification allows for "WeightClipping", which lets you set a maximum value (presumably absolute value) for the weights, but this does not enforce an equality constraint. (Aside: I suspect this is implemented with the MXNet clip_global_norm...and perhaps also relevant, is that the MXNet documentation doesn't seem to suggest a capability for applying the desired type of layer constraint.)

• One path would be to implement a custom network component using NetArrayLayer to store the learned weights and biases:

unitNormWeightLinearLayer[nIn_Integer, nOut_Integer] := NetGraph[
  <|"weight" -> NetArrayLayer["Output" -> {nOut, nIn}],
   "bias" -> NetArrayLayer["Output" -> {nOut}],
   "normalize" -> FunctionLayer[Map[Normalize]],
   "linear" -> FunctionLayer[ Apply[#1 . #2 + #3 &]]|>,
  {"weight" -> "normalize",
   {"normalize", NetPort["Input"], "bias"} -> 
    "linear" -> NetPort["Output"]},
  "Input" -> nIn]

unitNormWeightLinearLayer[5, 10]

• Although this only provides normalized values of the weights, it is perhaps suboptimal because it does not guarantee that the weights won't pathologically get too large during training (although perhaps that is doable by setting a "WeightClipping" to this layer to enforce the maximum value or attaching a loss function to the norm of the weight layer to keep it small) ... both of which seem a bit like hacky.

Is there a better way to implement this?