diff --git a/src/jaxfluids/solvers/riemann_solvers/LaxFriedrichs.py b/src/jaxfluids/solvers/riemann_solvers/LaxFriedrichs.py
index 916ecb9..f81a1a2 100644
--- a/src/jaxfluids/solvers/riemann_solvers/LaxFriedrichs.py
+++ b/src/jaxfluids/solvers/riemann_solvers/LaxFriedrichs.py
@@ -29,6 +29,8 @@
 #*                                                                              *
 #*------------------------------------------------------------------------------*
 
+import jax.numpy as jnp
+
 from jaxfluids.solvers.riemann_solvers.riemann_solver import RiemannSolver
 from jaxfluids.materials.material_manager import MaterialManager