Projected gradient for multidimensional array

Hi -- I have an array that is `mxnxn`. Each of the `m` `nxn` matrices represents a single "frame." I wish to enforce that each `nxn` matrix/frame is a doubly-stochastic matrix. For a single `nxn` matrix, I could use projected gradients with the [birkhoff](https://jaxopt.github.io/stable/_autosummary/jaxopt.projection.projection_birkhoff.html) projection. However, I would like to apply this constraint to *every* `nxn` matrix.

Does anyone know how to best implement this?