Run all examples in single precision

Author: pnkraemer

docs/examples_advanced/equinox_while_loop.py

@@ -27,9 +27,6 @@
 from probdiffeq import ivpsolve, ivpsolvers, taylor
 from probdiffeq.backend import control_flow
 
-jax.config.update("jax_platform_name", "cpu")
-
-
 # -
 
 # Overwrite the while-loop (via a context manager):

docs/examples_advanced/parameter_estimation_blackjax.py

@@ -135,11 +135,6 @@
 from probdiffeq import ivpsolve, ivpsolvers, stats, taylor
 
 # +
-# x64 precision
-jax.config.update("jax_enable_x64", True)
-
-# CPU
-jax.config.update("jax_platform_name", "cpu")
 
 # IVP examples in JAX
 if not backend.has_been_selected:

docs/examples_advanced/parameter_estimation_optax.py

@@ -37,8 +37,6 @@
 if not backend.has_been_selected:
     backend.select("jax")  # ivp examples in jax
 
-jax.config.update("jax_enable_x64", True)
-jax.config.update("jax_platform_name", "cpu")
 # -
 
 

docs/examples_basic/conditioning_on_zero_residual.py

@@ -32,9 +32,6 @@
 if not backend.has_been_selected:
     backend.select("jax")  # ivp examples in jax
 
-jax.config.update("jax_platform_name", "cpu")
-jax.config.update("jax_enable_x64", True)
-
 
 # +
 # Create an ODE problem

docs/examples_basic/dynamic_output_scales.py

@@ -46,7 +46,6 @@
     backend.select("jax")  # ivp examples in jax
 
 
-jax.config.update("jax_platform_name", "cpu")
 # -
 
 

docs/examples_basic/second_order_problems.py

@@ -28,7 +28,6 @@
 if not backend.has_been_selected:
     backend.select("jax")  # ivp examples in jax
 
-jax.config.update("jax_platform_name", "cpu")
 # -
 
 # Quick refresher: first-order ODEs

docs/examples_basic/taylor_coefficients.py

@@ -34,7 +34,6 @@
 if not backend.has_been_selected:
     backend.select("jax")  # ivp examples in jax
 
-jax.config.update("jax_platform_name", "cpu")
 # -
 
 # We start by defining an ODE.