flax

5edc65a8 - fix scale dtype and refactor q_dot_dq