diff --git a/frc971/control_loops/swerve/BUILD b/frc971/control_loops/swerve/BUILD
index 1b8dcc3..513bf1b 100644
--- a/frc971/control_loops/swerve/BUILD
+++ b/frc971/control_loops/swerve/BUILD
@@ -144,12 +144,14 @@
         "dynamics.cc",
         "dynamics.h",
         "dynamics.py",
+        "dynamics_constants.py",
     ],
     args = [
         "--output_base=$(BINDIR)/",
         "--cc_output_path=$(location :dynamics.cc)",
         "--h_output_path=$(location :dynamics.h)",
         "--casadi_py_output_path=$(location :dynamics.py)",
+        "--constants_output_path=$(location :dynamics_constants.py)",
     ],
     tool = ":generate_physics",
 )
@@ -204,10 +206,10 @@
 py_library(
     name = "jax_dynamics",
     srcs = [
+        "dynamics_constants.py",
         "jax_dynamics.py",
     ],
     deps = [
-        ":dynamics",
         "//frc971/control_loops/python:controls",
         "@pip//jax",
     ],
diff --git a/frc971/control_loops/swerve/generate_physics.cc b/frc971/control_loops/swerve/generate_physics.cc
index af058f6..7ea3db5 100644
--- a/frc971/control_loops/swerve/generate_physics.cc
+++ b/frc971/control_loops/swerve/generate_physics.cc
@@ -32,6 +32,8 @@
           "Path to write generated header code to");
 ABSL_FLAG(std::string, casadi_py_output_path, "",
           "Path to write casadi generated py code to");
+ABSL_FLAG(std::string, constants_output_path, "",
+          "Path to write constants python code to");
 ABSL_FLAG(double, caster, 0.01, "Caster in meters for the module.");
 
 ABSL_FLAG(bool, symbolic, false, "If true, write everything out symbolically.");
@@ -625,6 +627,80 @@
     }
   }
 
+  void WriteConstantsFile(std::string_view path) {
+    std::vector<std::string> result_py;
+
+    // Write out the header.
+    result_py.emplace_back("#!/usr/bin/env python3");
+    result_py.emplace_back("");
+
+    WriteConstants(&result_py);
+
+    aos::util::WriteStringToFileOrDie(path, absl::StrJoin(result_py, "\n"));
+  }
+
+  void WriteConstants(std::vector<std::string> *result_py) {
+    result_py->emplace_back(absl::Substitute("WHEEL_RADIUS = $0", ccode(*rw_)));
+    result_py->emplace_back(
+        absl::Substitute("ROBOT_WIDTH = $0", ccode(*robot_width_)));
+    result_py->emplace_back(absl::Substitute("CASTER = $0", ccode(*caster_)));
+    result_py->emplace_back("STATE_THETAS0 = 0");
+    result_py->emplace_back("STATE_THETAD0 = 1");
+    result_py->emplace_back("STATE_OMEGAS0 = 2");
+    result_py->emplace_back("STATE_OMEGAD0 = 3");
+    result_py->emplace_back("STATE_THETAS1 = 4");
+    result_py->emplace_back("STATE_THETAD1 = 5");
+    result_py->emplace_back("STATE_OMEGAS1 = 6");
+    result_py->emplace_back("STATE_OMEGAD1 = 7");
+    result_py->emplace_back("STATE_THETAS2 = 8");
+    result_py->emplace_back("STATE_THETAD2 = 9");
+    result_py->emplace_back("STATE_OMEGAS2 = 10");
+    result_py->emplace_back("STATE_OMEGAD2 = 11");
+    result_py->emplace_back("STATE_THETAS3 = 12");
+    result_py->emplace_back("STATE_THETAD3 = 13");
+    result_py->emplace_back("STATE_OMEGAS3 = 14");
+    result_py->emplace_back("STATE_OMEGAD3 = 15");
+    result_py->emplace_back("STATE_X = 16");
+    result_py->emplace_back("STATE_Y = 17");
+    result_py->emplace_back("STATE_THETA = 18");
+    result_py->emplace_back("STATE_VX = 19");
+    result_py->emplace_back("STATE_VY = 20");
+    result_py->emplace_back("STATE_OMEGA = 21");
+    result_py->emplace_back("STATE_FX = 22");
+    result_py->emplace_back("STATE_FY = 23");
+    result_py->emplace_back("STATE_MOMENT = 24");
+    result_py->emplace_back("NUM_STATES = 25");
+    result_py->emplace_back("");
+    result_py->emplace_back("VELOCITY_STATE_THETAS0 = 0");
+    result_py->emplace_back("VELOCITY_STATE_OMEGAS0 = 1");
+    result_py->emplace_back("VELOCITY_STATE_THETAS1 = 2");
+    result_py->emplace_back("VELOCITY_STATE_OMEGAS1 = 3");
+    result_py->emplace_back("VELOCITY_STATE_THETAS2 = 4");
+    result_py->emplace_back("VELOCITY_STATE_OMEGAS2 = 5");
+    result_py->emplace_back("VELOCITY_STATE_THETAS3 = 6");
+    result_py->emplace_back("VELOCITY_STATE_OMEGAS3 = 7");
+    result_py->emplace_back("VELOCITY_STATE_THETA = 8");
+    result_py->emplace_back("VELOCITY_STATE_VX = 9");
+    result_py->emplace_back("VELOCITY_STATE_VY = 10");
+    result_py->emplace_back("VELOCITY_STATE_OMEGA = 11");
+    // result_py->emplace_back("VELOCITY_STATE_FX = 16");
+    // result_py->emplace_back("VELOCITY_STATE_FY = 17");
+    // result_py->emplace_back("VELOCITY_STATE_MOMENT = 18");
+    result_py->emplace_back("NUM_VELOCITY_STATES = 12");
+    result_py->emplace_back("");
+    result_py->emplace_back("");
+    result_py->emplace_back("# Is = STEER_CURRENT_COUPLING_FACTOR * Id");
+    result_py->emplace_back(absl::Substitute(
+        "STEER_CURRENT_COUPLING_FACTOR = $0",
+        ccode(*(neg(
+            mul(div(Gs_, Kts_),
+                mul(div(Ktd_, mul(Gd_, rw_)),
+                    neg(mul(add(neg(wb_), mul(add(rs_, rp_),
+                                              sub(integer(1), div(rb1_, rp_)))),
+                            div(rw_, rb2_))))))))));
+    result_py->emplace_back("");
+  }
+
   // Writes the physics out to the provided .cc and .h path.
   void WriteCasadi(std::string_view py_path) {
     std::vector<std::string> result_py;
@@ -634,54 +710,9 @@
     result_py.emplace_back("");
     result_py.emplace_back("import casadi, numpy");
     result_py.emplace_back("");
-    result_py.emplace_back(absl::Substitute("WHEEL_RADIUS = $0", ccode(*rw_)));
-    result_py.emplace_back(
-        absl::Substitute("ROBOT_WIDTH = $0", ccode(*robot_width_)));
-    result_py.emplace_back(absl::Substitute("CASTER = $0", ccode(*caster_)));
-    result_py.emplace_back("STATE_THETAS0 = 0");
-    result_py.emplace_back("STATE_THETAD0 = 1");
-    result_py.emplace_back("STATE_OMEGAS0 = 2");
-    result_py.emplace_back("STATE_OMEGAD0 = 3");
-    result_py.emplace_back("STATE_THETAS1 = 4");
-    result_py.emplace_back("STATE_THETAD1 = 5");
-    result_py.emplace_back("STATE_OMEGAS1 = 6");
-    result_py.emplace_back("STATE_OMEGAD1 = 7");
-    result_py.emplace_back("STATE_THETAS2 = 8");
-    result_py.emplace_back("STATE_THETAD2 = 9");
-    result_py.emplace_back("STATE_OMEGAS2 = 10");
-    result_py.emplace_back("STATE_OMEGAD2 = 11");
-    result_py.emplace_back("STATE_THETAS3 = 12");
-    result_py.emplace_back("STATE_THETAD3 = 13");
-    result_py.emplace_back("STATE_OMEGAS3 = 14");
-    result_py.emplace_back("STATE_OMEGAD3 = 15");
-    result_py.emplace_back("STATE_X = 16");
-    result_py.emplace_back("STATE_Y = 17");
-    result_py.emplace_back("STATE_THETA = 18");
-    result_py.emplace_back("STATE_VX = 19");
-    result_py.emplace_back("STATE_VY = 20");
-    result_py.emplace_back("STATE_OMEGA = 21");
-    result_py.emplace_back("STATE_FX = 22");
-    result_py.emplace_back("STATE_FY = 23");
-    result_py.emplace_back("STATE_MOMENT = 24");
-    result_py.emplace_back("NUM_STATES = 25");
-    result_py.emplace_back("");
-    result_py.emplace_back("VELOCITY_STATE_THETAS0 = 0");
-    result_py.emplace_back("VELOCITY_STATE_OMEGAS0 = 1");
-    result_py.emplace_back("VELOCITY_STATE_THETAS1 = 2");
-    result_py.emplace_back("VELOCITY_STATE_OMEGAS1 = 3");
-    result_py.emplace_back("VELOCITY_STATE_THETAS2 = 4");
-    result_py.emplace_back("VELOCITY_STATE_OMEGAS2 = 5");
-    result_py.emplace_back("VELOCITY_STATE_THETAS3 = 6");
-    result_py.emplace_back("VELOCITY_STATE_OMEGAS3 = 7");
-    result_py.emplace_back("VELOCITY_STATE_THETA = 8");
-    result_py.emplace_back("VELOCITY_STATE_VX = 9");
-    result_py.emplace_back("VELOCITY_STATE_VY = 10");
-    result_py.emplace_back("VELOCITY_STATE_OMEGA = 11");
-    // result_py.emplace_back("VELOCITY_STATE_FX = 16");
-    // result_py.emplace_back("VELOCITY_STATE_FY = 17");
-    // result_py.emplace_back("VELOCITY_STATE_MOMENT = 18");
-    result_py.emplace_back("NUM_VELOCITY_STATES = 12");
-    result_py.emplace_back("");
+
+    WriteConstants(&result_py);
+
     result_py.emplace_back("def to_velocity_state(X):");
     result_py.emplace_back("    return numpy.array([");
     result_py.emplace_back("        [X[STATE_THETAS0, 0]],");
@@ -732,27 +763,6 @@
                            "casadi.exp($1.0 * x))) * $0.0]))) / $0.0)",
                            kLogGain, kAbsGain));
     }
-    result_py.emplace_back("");
-    result_py.emplace_back("# Is = STEER_CURRENT_COUPLING_FACTOR * Id");
-    result_py.emplace_back(absl::Substitute(
-        "STEER_CURRENT_COUPLING_FACTOR = $0",
-        ccode(*(neg(
-            mul(div(Gs_, Kts_),
-                mul(div(Ktd_, mul(Gd_, rw_)),
-                    neg(mul(add(neg(wb_), mul(add(rs_, rp_),
-                                              sub(integer(1), div(rb1_, rp_)))),
-                            div(rw_, rb2_))))))))));
-    result_py.emplace_back("");
-    result_py.emplace_back("# Is = STEER_CURRENT_COUPLING_FACTOR * Id");
-    result_py.emplace_back(absl::Substitute(
-        "STEER_CURRENT_COUPLING_FACTOR = $0",
-        ccode(*(neg(
-            mul(div(Gs_, Kts_),
-                mul(div(Ktd_, mul(Gd_, rw_)),
-                    neg(mul(add(neg(wb_), mul(add(rs_, rp_),
-                                              sub(integer(1), div(rb1_, rp_)))),
-                            div(rw_, rb2_))))))))));
-    result_py.emplace_back("");
 
     result_py.emplace_back("# Returns the derivative of our state vector");
     result_py.emplace_back("# [thetas0, thetad0, omegas0, omegad0,");
@@ -1284,5 +1294,9 @@
     sim.WriteCasadi(absl::GetFlag(FLAGS_casadi_py_output_path));
   }
 
+  if (!absl::GetFlag(FLAGS_constants_output_path).empty()) {
+    sim.WriteConstantsFile(absl::GetFlag(FLAGS_constants_output_path));
+  }
+
   return 0;
 }
diff --git a/frc971/control_loops/swerve/jax_dynamics.py b/frc971/control_loops/swerve/jax_dynamics.py
index dfadbda..6e1305d 100644
--- a/frc971/control_loops/swerve/jax_dynamics.py
+++ b/frc971/control_loops/swerve/jax_dynamics.py
@@ -4,8 +4,8 @@
 from collections import namedtuple
 import jax
 
-from frc971.control_loops.swerve import dynamics
 from frc971.control_loops.python.control_loop import KrakenFOC
+from frc971.control_loops.swerve.dynamics_constants import *
 
 # Note: this physics needs to match the symengine code.  We have tests that
 # confirm they match in all the cases we care about.
@@ -130,36 +130,35 @@
             jax.numpy.array([1.0, softabs_x * kMaxLogGain])) / kMaxLogGain)
 
 
-def full_module_physics(coefficients: dict, Rtheta, module_index: int,
-                        mounting_location, X, U):
+def full_module_physics(coefficients: CoefficientsType, Rtheta,
+                        module_index: int, mounting_location, X, U):
     X_module = X[module_index * 4:(module_index + 1) * 4]
     Is = U[2 * module_index + 0]
     Id = U[2 * module_index + 1]
 
-    Rthetaplusthetas = R(X[dynamics.STATE_THETA] +
-                         X_module[dynamics.STATE_THETAS0])
+    Rthetaplusthetas = R(X[STATE_THETA] + X_module[STATE_THETAS0])
 
     caster_vector = jax.numpy.array([-coefficients.caster, 0.0])
 
-    robot_velocity = X[dynamics.STATE_VX:dynamics.STATE_VY + 1]
+    robot_velocity = X[STATE_VX:STATE_VY + 1]
 
     contact_patch_velocity = (
-        angle_cross(Rtheta @ mounting_location, X[dynamics.STATE_OMEGA]) +
-        robot_velocity + angle_cross(
-            Rthetaplusthetas @ caster_vector,
-            (X[dynamics.STATE_OMEGA] + X_module[dynamics.STATE_OMEGAS0])))
+        angle_cross(Rtheta @ mounting_location, X[STATE_OMEGA]) +
+        robot_velocity +
+        angle_cross(Rthetaplusthetas @ caster_vector,
+                    (X[STATE_OMEGA] + X_module[STATE_OMEGAS0])))
 
     wheel_ground_velocity = Rthetaplusthetas.T @ contact_patch_velocity
 
     wheel_velocity = jax.numpy.array(
-        [coefficients.rw * X_module[dynamics.STATE_OMEGAD0], 0.0])
+        [coefficients.rw * X_module[STATE_OMEGAD0], 0.0])
 
     wheel_slip_velocity = wheel_velocity - wheel_ground_velocity
 
     slip_angle = jax.numpy.sin(
         -soft_atan2(wheel_ground_velocity[1], wheel_ground_velocity[0]))
 
-    slip_ratio = (coefficients.rw * X_module[dynamics.STATE_OMEGAD0] -
+    slip_ratio = (coefficients.rw * X_module[STATE_OMEGAD0] -
                   wheel_ground_velocity[0]) / jax.numpy.max(
                       jax.numpy.array(
                           [0.02, jax.numpy.abs(wheel_ground_velocity[0])]))
@@ -193,8 +192,8 @@
 
     X_dot_contribution = jax.numpy.hstack((jax.numpy.zeros(
         (4, )), ) * (module_index) + (jax.numpy.array([
-            X_module[dynamics.STATE_OMEGAS0],
-            X_module[dynamics.STATE_OMEGAD0],
+            X_module[STATE_OMEGAS0],
+            X_module[STATE_OMEGAD0],
             alphas,
             alphad,
         ]), ) + (jax.numpy.zeros((4, )), ) * (3 - module_index) + (
@@ -208,7 +207,7 @@
 
 @partial(jax.jit, static_argnames=['coefficients'])
 def full_dynamics(coefficients: CoefficientsType, X, U):
-    Rtheta = R(X[dynamics.STATE_THETA])
+    Rtheta = R(X[STATE_THETA])
 
     module0 = full_module_physics(
         coefficients, Rtheta, 0,
@@ -233,37 +232,41 @@
 
     X_dot = module0 + module1 + module2 + module3
 
-    X_dot = X_dot.at[dynamics.STATE_X:dynamics.STATE_THETA + 1].set(
+    X_dot = X_dot.at[STATE_X:STATE_THETA + 1].set(
         jax.numpy.array([
-            X[dynamics.STATE_VX],
-            X[dynamics.STATE_VY],
-            X[dynamics.STATE_OMEGA],
+            X[STATE_VX],
+            X[STATE_VY],
+            X[STATE_OMEGA],
         ]))
 
     return X_dot
 
 
-def velocity_module_physics(coefficients: dict, Rtheta, module_index: int,
-                            mounting_location, X, U):
+def velocity_module_physics(coefficients: CoefficientsType,
+                            Rtheta: jax.typing.ArrayLike, module_index: int,
+                            mounting_location: jax.typing.ArrayLike,
+                            X: jax.typing.ArrayLike, U: jax.typing.ArrayLike):
     X_module = X[module_index * 2:(module_index + 1) * 2]
     Is = U[2 * module_index + 0]
     Id = U[2 * module_index + 1]
 
-    Rthetaplusthetas = R(X[dynamics.VELOCITY_STATE_THETA] +
-                         X_module[dynamics.VELOCITY_STATE_THETAS0])
+    rotated_mounting_location = Rtheta @ mounting_location
+
+    Rthetaplusthetas = R(X[VELOCITY_STATE_THETA] +
+                         X_module[VELOCITY_STATE_THETAS0])
 
     caster_vector = jax.numpy.array([-coefficients.caster, 0.0])
 
-    robot_velocity = X[dynamics.VELOCITY_STATE_VX:dynamics.VELOCITY_STATE_VY +
-                       1]
+    robot_velocity = X[VELOCITY_STATE_VX:VELOCITY_STATE_VY + 1]
 
     contact_patch_velocity = (
-        angle_cross(Rtheta @ mounting_location,
-                    X[dynamics.VELOCITY_STATE_OMEGA]) + robot_velocity +
-        angle_cross(Rthetaplusthetas @ caster_vector,
-                    (X[dynamics.VELOCITY_STATE_OMEGA] +
-                     X_module[dynamics.VELOCITY_STATE_OMEGAS0])))
+        angle_cross(rotated_mounting_location, X[VELOCITY_STATE_OMEGA]) +
+        robot_velocity + angle_cross(
+            Rthetaplusthetas @ caster_vector,
+            (X[VELOCITY_STATE_OMEGA] + X_module[VELOCITY_STATE_OMEGAS0])))
 
+    # Velocity of the contact patch over the field projected into the direction
+    # of the wheel.
     wheel_ground_velocity = Rthetaplusthetas.T @ contact_patch_velocity
 
     slip_angle = jax.numpy.sin(
@@ -288,11 +291,11 @@
 
     F = Rthetaplusthetas @ jax.numpy.array([Fwx, Fwy])
 
-    torque = force_cross(Rtheta @ mounting_location, F)
+    torque = force_cross(rotated_mounting_location, F)
 
     X_dot_contribution = jax.numpy.hstack((jax.numpy.zeros(
         (2, )), ) * (module_index) + (jax.numpy.array([
-            X_module[dynamics.VELOCITY_STATE_OMEGAS0],
+            X_module[VELOCITY_STATE_OMEGAS0],
             alphas,
         ]), ) + (jax.numpy.zeros((2, )), ) * (3 - module_index) + (
             jax.numpy.zeros((1, )),
@@ -300,29 +303,30 @@
             jax.numpy.array([torque / coefficients.J]),
         ))
 
-    return X_dot_contribution
+    return X_dot_contribution, F, torque
 
 
 @partial(jax.jit, static_argnames=['coefficients'])
-def velocity_dynamics(coefficients: CoefficientsType, X, U):
-    Rtheta = R(X[dynamics.VELOCITY_STATE_THETA])
+def velocity_dynamics(coefficients: CoefficientsType, X: jax.typing.ArrayLike,
+                      U: jax.typing.ArrayLike):
+    Rtheta = R(X[VELOCITY_STATE_THETA])
 
-    module0 = velocity_module_physics(
+    module0, _, _ = velocity_module_physics(
         coefficients, Rtheta, 0,
         jax.numpy.array(
             [coefficients.robot_width / 2.0, coefficients.robot_width / 2.0]),
         X, U)
-    module1 = velocity_module_physics(
+    module1, _, _ = velocity_module_physics(
         coefficients, Rtheta, 1,
         jax.numpy.array(
             [-coefficients.robot_width / 2.0, coefficients.robot_width / 2.0]),
         X, U)
-    module2 = velocity_module_physics(
+    module2, _, _ = velocity_module_physics(
         coefficients, Rtheta, 2,
         jax.numpy.array(
             [-coefficients.robot_width / 2.0,
              -coefficients.robot_width / 2.0]), X, U)
-    module3 = velocity_module_physics(
+    module3, _, _ = velocity_module_physics(
         coefficients, Rtheta, 3,
         jax.numpy.array(
             [coefficients.robot_width / 2.0, -coefficients.robot_width / 2.0]),
@@ -330,5 +334,21 @@
 
     X_dot = module0 + module1 + module2 + module3
 
-    return X_dot.at[dynamics.VELOCITY_STATE_THETA].set(
-        X[dynamics.VELOCITY_STATE_OMEGA])
+    return X_dot.at[VELOCITY_STATE_THETA].set(X[VELOCITY_STATE_OMEGA])
+
+
+def to_velocity_state(X):
+    return jax.numpy.array([
+        X[STATE_THETAS0],
+        X[STATE_OMEGAS0],
+        X[STATE_THETAS1],
+        X[STATE_OMEGAS1],
+        X[STATE_THETAS2],
+        X[STATE_OMEGAS2],
+        X[STATE_THETAS3],
+        X[STATE_OMEGAS3],
+        X[STATE_THETA],
+        X[STATE_VX],
+        X[STATE_VY],
+        X[STATE_OMEGA],
+    ])