How can you constrain a neural network layer to simply be a n-dimensional rotation layer? I'm looking to constrain one layer of my neural network to specifically find the best rotation of its input in order to satisfy an objective.  (My end goal, where $R$ is the rotation layer, is of the form $R^T ~f_{objective} \left(Rz \right)$  ).
I am looking to train this (+ other components) via gradient descent. If $z\in\mathbb{R}^2$, then I can just say
$R = \left[ \begin{matrix}
\cos \theta & -\sin \theta\\
\sin \theta & \cos \theta
\end{matrix} \right] $,
and have $\theta$ be a learnable parameter.
However, I am lost on how to actually set this up for an $d$-dimensional space (where $d$>10). I've tried looking at resources on how to make a $d$-dimensional rotation matrix and it gets heavy into Linear Algebra and is way over my head. It feels like this should be easier than it seems, so I feel like I'm overlooking something (like maybe $R$ should just be a usual linear layer without any non-linear activations).
Anyone have any ideas? I appreciate you, in advance : )

The most computationally efficient representation of a rotation is as a series of $n$ householder reflections. The concept is pretty simple: your rotation is parameterized by a series of $n$ arbitrary vectors  $v_1\dots v_n$ each with $n$ elements. To apply a corresponding rotation to an input $x$ vector, iterate through each of your $n$ parameter vectors $v_i$ and reflect $x$ through that vector
$x \leftarrow x - 2v_i\frac{x \cdot v_i}{v_i \cdot v_i}$
Alternatively, you can construct your entire rotation matrix explicitely from $R = \overset{n}{\underset{i=1}{\prod}}\left(I - 2 \frac{v_i \otimes v_i}{v_i \cdot v_i}\right)$
See here for a publication on this. Autograd frameworks like pytorch, tensorflow, or jax will have no problem backpropagating through this transformation so you can optimize parameters $v$.
