Trajectory Optimization in JAX A toolbox for trajectory optimization and model-predictive control based on JAX