@@ -0,0 +1,2028 @@
+# Copyright 2020 Google LLC%0A#%0A# Licensed under the Apache License, Version 2.0 (the %22License%22);%0A# you may not use this file except in compliance with the License.%0A# You may obtain a copy of the License at%0A#%0A#     https://www.apache.org/licenses/LICENSE-2.0%0A#%0A# Unless required by applicable law or agreed to in writing, software%0A# distributed under the License is distributed on an %22AS IS%22 BASIS,%0A# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.%0A# See the License for the specific language governing permissions and%0A# limitations under the License.%0A%0Afrom functools import partial%0A%0Afrom absl.testing import absltest%0Aimport numpy as onp%0A%0Afrom jax import jit%0Afrom jax import dtypes%0Afrom jax import test_util as jtu%0Aimport jax.numpy as np%0Afrom jax.experimental.ode import odeint%0A%0Afrom jax.config import config%0Aconfig.parse_flags_with_absl()%0A%0A%0Adef num_float_bits(dtype):%0A  return dtypes.finfo(dtypes.canonicalize_dtype(dtype)).bits%0A%0A%0Aclass JetTest(jtu.JaxTestCase):%0A%0A  @jtu.skip_on_devices(%22tpu%22)%0A  def test_pend_grads(self):%0A    def pend(y, _, m, g):%0A      theta, omega = y%0A      return %5Bomega, -m * omega - g * np.sin(theta)%5D%0A%0A    def f(y0, ts, *args):%0A      return odeint(pend, y0, ts, *args)%0A%0A    y0 = %5Bnp.pi - 0.1, 0.0%5D%0A    ts = np.linspace(0., 1., 11)%0A    args = (0.25, 9.8)%0A%0A    tol = 1e-1 if num_float_bits(onp.float64) == 32 else 1e-3%0A    jtu.check_grads(f, (y0, ts, *args), modes=%5B%22rev%22%5D, order=2,%0A                    atol=tol, rtol=tol)%0A%0A  @jtu.skip_on_devices(%22tpu%22)%0A  def test_weird_time_pendulum_grads(self):%0A    %22%22%22Test that gradients are correct when the dynamics depend on t.%22%22%22%0A    def dynamics(y, t):%0A      return np.array(%5By%5B1%5D * -t, -1 * y%5B1%5D - 9.8 * np.sin(y%5B0%5D)%5D)%0A%0A    integrate = partial(odeint, dynamics)%0A%0A    y0 = %5Bnp.pi - 0.1, 0.0%5D%0A    ts = np.linspace(0., 1., 11)%0A%0A    tol = 1e-1 if num_float_bits(onp.float64) == 32 else 1e-3%0A    jtu.check_grads(integrate, (y0, ts), modes=%5B%22rev%22%5D, order=2,%0A                    rtol=tol, atol=tol)%0A%0A%0Aif __name__ == '__main__':%0A  absltest.main()%0A

add ode test file (#2818)