UnitNorm constraint

Is there an existing pytorch.distribution.constraint for unit norms of the innermost dimension? The simplex constraint is quite nice - would be great to have something just like that. Seems like Keras has a UnitNorm constraint. Might be a nice addition to the core set.

Thanks!