I was struggling to maximize Tr(W'X'XW) with constraint W'W = I in TensorFlow. The problematic part is the constraint since it requires the weights stay on a Stiefel manifold. There is a soft way which adds the constraint as a Lagrangian term to the objective. However, this is not exact, and implementation can be messy.
I'm aware of the Pymanopt library that uses autodiff for manifold optimization. But I am still wondering if there was a principle way to simulate (Stiefel) manifold optimization within Tensorflow?