Add regression test

jax-ml · Apr 24, 2020 · 04bae57 · 04bae57
1 parent c7d597f
commit 04bae57
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/jax/experimental/ode.py b/jax/experimental/ode.py
@@ -295,7 +295,17 @@ def f(y0, ts, *args):
   check_grads(f, (y0, ts, *args), modes=["rev"], order=2,
               atol=1e-1, rtol=1e-1)
 
+def weird_time_pendulum_check_grads():
+  """Test that gradients are correct when the dynamics depend on t."""
+  def f(y0, ts):
+    return odeint(lambda y, t: np.array([y[1] * -t, -1 * y[1] - 9.8 * np.sin(y[0])]), y0, ts)
+
+  y0 = [np.pi - 0.1, 0.0]
+  ts = np.linspace(0., 1., 11)
+
+  check_grads(f, (y0, ts), modes=["rev"], order=2)
 
 if __name__ == '__main__':
   pend_benchmark_odeint()
   pend_check_grads()
+  weird_time_pendulum_check_grads()