diff --git a/.github/workflows/install_and_test.yml b/.github/workflows/install_and_test.yml
index 37c0c22..590b6ee 100644
--- a/.github/workflows/install_and_test.yml
+++ b/.github/workflows/install_and_test.yml
@@ -14,7 +14,7 @@ jobs:
       PYSCF_CONFIG_FILE: ".github/workflows/pyscf_conf.py"
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11"]
+        python-version: ["3.12"]
         os: [ubuntu-latest]
     steps:
       - uses: actions/checkout@v3
@@ -32,8 +32,13 @@ jobs:
           pip install -e ".[examples]"
       - name: Run unit tests
         run: |
+            pytest -v tests/unit/test_kinetic_density_eq.py
             pytest -v tests/unit/test_eigenproblem.py
             pytest -v tests/unit/test_loss.py
+      - name: Debug
+        run: |
+            pip list
+            python --version
       - name: Run integration tests
         run: |
           pytest -v tests/integration/molecules/test_non_xc_energy.py
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py
index 7c62fa2..811b92a 100644
--- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py
+++ b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py
@@ -216,7 +216,7 @@ def _build_graph(self, batch_dim: Optional[int] = None):
             library.
         """
 
-        self._functional = hub.Module(spec=self._model_path)
+        self._functional = hub.KerasLayer(self._model_path, signature_outputs_as_dict=True)
 
         grid_coords = tf.placeholder(tf.float32, shape=[batch_dim, 3], name="grid_coords")
         grid_weights = tf.placeholder(tf.float32, shape=[batch_dim], name="grid_weights")
@@ -260,7 +260,7 @@ def _build_graph(self, batch_dim: Optional[int] = None):
         }
         tensor_dict = {f"tensor_dict${k}": v for k, v in features.items()}
 
-        predictions = self._functional(tensor_dict, as_dict=True)
+        predictions = self._functional(tensor_dict)
         local_xc = predictions["grid_contribution"]
         weighted_local_xc = local_xc * grid_weights
         unweighted_xc = tf.reduce_sum(local_xc, axis=0)
@@ -311,19 +311,6 @@ def _build_graph(self, batch_dim: Optional[int] = None):
             grid_weights=grid_weights,
         )
 
-        outputs = {
-            "vxc": self._vxc,
-            "vrho": tf.stack(self._vrho),
-            "vsigma": tf.stack(self._vsigma),
-            "vtau": tf.stack(self._vtau),
-            "vhf": tf.stack(self._vhf),
-        }
-        # Create the signature for TF-Hub, including both the energy and functional
-        # derivatives.
-        # This is a no-op if _build_graph is called outside of
-        # hub.create_module_spec.
-        hub.add_signature(inputs=attr.asdict(self._placeholders), outputs=outputs)
-
     def export_functional_and_derivatives(
         self,
         export_path: str,
@@ -351,7 +338,7 @@ def export_functional_and_derivatives(
             spec = hub.create_module_spec(
                 self._build_graph, tags_and_args=[(set(), {"batch_dim": batch_dim})]
             )
-            functional_and_derivatives = hub.Module(spec=spec)
+            functional_and_derivatives = hub.KerasLayer(spec, signature_outputs_as_dict=True)
             with tf.Session() as session:
                 session.run(tf.global_variables_initializer())
                 functional_and_derivatives.export(export_path, session)
diff --git a/grad_dft/functional.py b/grad_dft/functional.py
index 6e5ffef..e55a00d 100644
--- a/grad_dft/functional.py
+++ b/grad_dft/functional.py
@@ -591,10 +591,10 @@ def dm21_densities(
     grad_rho_norm_sq = jnp.sum(grad_rho**2, axis=-1)
 
     # LDA preprocessing data
-    log_rho = jnp.log2(jnp.clip(rho, a_min=clip_cte))
+    log_rho = jnp.log2(jnp.clip(rho, clip_cte))
 
     # GGA preprocessing data
-    log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, a_min=clip_cte)) / 2
+    log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, clip_cte)) / 2
     log_x_sigma = log_grad_rho_norm - 4 / 3.0 * log_rho
     log_u_sigma = jnp.where(
         jnp.greater(log_rho, jnp.log2(clip_cte)),
@@ -603,7 +603,7 @@ def dm21_densities(
     )
 
     # MGGA preprocessing data
-    log_tau = jnp.log2(jnp.clip(tau, a_min=clip_cte))
+    log_tau = jnp.log2(jnp.clip(tau, clip_cte))
     log_1t_sigma = -(
         5 / 3.0 * log_rho - log_tau + 2 / 3.0 * jnp.log2(6 * jnp.pi**2) + jnp.log2(3 / 5.0)
     )
@@ -674,8 +674,7 @@ def dm21_combine_densities(
         [densities] + [ehf[i].sum(axis=0, keepdims=True).T for i in range(len(ehf))], axis=1
     )
 
-@jaxtyped
-@typechecked
+@jaxtyped(typechecker=typechecked)
 def dm21_hfgrads_densities(
     functional: nn.Module,
     params: PyTree,
@@ -716,8 +715,7 @@ def dm21_hfgrads_densities(
     )
     return vxc_hf.sum(axis=0)  # Sum over omega
 
-@jaxtyped
-@typechecked
+@jaxtyped(typechecker=typechecked)
 def dm21_hfgrads_cinputs(
     functional: nn.Module,
     params: PyTree,
@@ -1005,7 +1003,7 @@ def correlation_polarization_correction(
         The ready to be integrated electronic energy density.
     """
 
-    log_rho = jnp.log2(jnp.clip(rho.sum(axis=1), a_min=clip_cte))
+    log_rho = jnp.log2(jnp.clip(rho.sum(axis=1), clip_cte))
     # assert not jnp.isnan(log_rho).any() and not jnp.isinf(log_rho).any()
     log_rs = jnp.log2((3 / (4 * jnp.pi)) ** (1 / 3)) - log_rho / 3.0
 
@@ -1033,7 +1031,7 @@ def fzeta(z):
     # assert not jnp.isnan(alphac).any() and not jnp.isinf(alphac).any()
 
     fz = fzeta(zeta) #jnp.round(fzeta(zeta), int(math.log10(clip_cte)))
-    z4 = zeta**4 #jnp.round(2 ** (4 * jnp.log2(jnp.clip(zeta, a_min=clip_cte))), int(math.log10(clip_cte)))
+    z4 = zeta**4 #jnp.round(2 ** (4 * jnp.log2(jnp.clip(zeta, clip_cte))), int(math.log10(clip_cte)))
 
     e_tilde = (
         e_tilde_PF[:, 0]
@@ -1105,10 +1103,10 @@ def densities(
     grad_rho_norm_sq = jnp.sum(grad_rho**2, axis=-1)
 
     # LDA preprocessing data
-    log_rho = jnp.log2(jnp.clip(rho, a_min=clip_cte))
+    log_rho = jnp.log2(jnp.clip(rho, clip_cte))
 
     # GGA preprocessing data
-    log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, a_min=clip_cte)) / 2
+    log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, clip_cte)) / 2
     log_x_sigma = log_grad_rho_norm - 4 / 3.0 * log_rho
     log_u_sigma = jnp.where(
         jnp.greater(log_rho, jnp.log2(clip_cte)),
@@ -1117,7 +1115,7 @@ def densities(
     )
 
     # MGGA preprocessing data
-    log_tau = jnp.log2(jnp.clip(tau, a_min=clip_cte))
+    log_tau = jnp.log2(jnp.clip(tau, clip_cte))
     log_1t_sigma = log_tau - 5 / 3.0 * log_rho
     log_w_sigma = jnp.where(
         jnp.greater(log_rho, jnp.log2(clip_cte)),
@@ -1136,8 +1134,8 @@ def densities(
     ######### Correlation features ###############
 
     grad_rho_norm_sq_ss = jnp.sum((grad_rho.sum(axis=1)) ** 2, axis=-1)
-    log_grad_rho_norm_ss = jnp.log2(jnp.clip(grad_rho_norm_sq_ss, a_min=clip_cte)) / 2
-    log_rho_ss = jnp.log2(jnp.clip(rho.sum(axis=1), a_min=clip_cte))
+    log_grad_rho_norm_ss = jnp.log2(jnp.clip(grad_rho_norm_sq_ss, clip_cte)) / 2
+    log_rho_ss = jnp.log2(jnp.clip(rho.sum(axis=1), clip_cte))
     log_x_ss = log_grad_rho_norm_ss - 4 / 3.0 * log_rho_ss
 
     log_u_ss = jnp.where(
@@ -1154,7 +1152,7 @@ def densities(
 
     log_u_c = jnp.stack((log_u_ss, log_u_ab), axis=1)
 
-    log_tau_ss = jnp.log2(jnp.clip(tau.sum(axis=1), a_min=clip_cte))
+    log_tau_ss = jnp.log2(jnp.clip(tau.sum(axis=1), clip_cte))
     log_1t_ss = log_tau_ss - 5 / 3.0 * log_rho_ss
     log_w_ss = jnp.where(
         jnp.greater(log_rho.sum(axis=1), jnp.log2(clip_cte)),
@@ -1177,7 +1175,7 @@ def densities(
     beta3 = jnp.array([[1.6382, 3.3662]])
     beta4 = jnp.array([[0.49294, 0.62517]])
 
-    log_rho = jnp.log2(jnp.clip(rho.sum(axis=1, keepdims=True), a_min=clip_cte))
+    log_rho = jnp.log2(jnp.clip(rho.sum(axis=1, keepdims=True), clip_cte))
     log_rs = jnp.log2((3 / (4 * jnp.pi)) ** (1 / 3)) - log_rho / 3.0
     brs_1_2 = 2 ** (log_rs / 2 + jnp.log2(beta1))
     ars = 2 ** (log_rs + jnp.log2(alpha1))
diff --git a/grad_dft/molecule.py b/grad_dft/molecule.py
index 580430b..a0ef2f1 100644
--- a/grad_dft/molecule.py
+++ b/grad_dft/molecule.py
@@ -338,9 +338,8 @@ def to_dict(self) -> dict:
 
 #######################################################################
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def orbital_grad(
         mo_coeff: Float[Array, "spin orbitals orbitals"],
         mo_occ: Float[Array, "spin orbitals"],
@@ -382,9 +381,8 @@ def orbital_grad(
 
 
 ##########################################################
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def density(rdm1: Float[Array, "spin orbitals orbitals"], 
             ao: Float[Array, "grid orbitals"], 
             precision: Precision = Precision.HIGHEST
@@ -408,9 +406,8 @@ def density(rdm1: Float[Array, "spin orbitals orbitals"],
 
     return jnp.einsum("...ab,ra,rb->r...", rdm1, ao, ao, precision=precision)
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def grad_density(
     rdm1: Float[Array, "spin orbitals orbitals"], 
     ao: Float[Array, "grid orbitals"], 
@@ -439,9 +436,8 @@ def grad_density(
 
     return 2 * jnp.einsum("...ab,ra,rbj->r...j", rdm1, ao, grad_ao, precision=precision)
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def lapl_density(
     rdm1: Float[Array, "spin orbitals orbitals"], 
     ao: Float[Array, "grid orbitals"], 
@@ -473,9 +469,8 @@ def lapl_density(
         "...ab,raj,rbj->r...", rdm1, grad_ao, grad_ao, precision=precision
     ) + 2 * jnp.einsum("...ab,ra,rbi->r...", rdm1, ao, grad_2_ao, precision=precision)
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def kinetic_density(
     rdm1: Float[Array, "spin orbitals orbitals"], 
     grad_ao: Float[Array, "grid orbitals 3"],
@@ -498,12 +493,11 @@ def kinetic_density(
     Array
         The kinetic energy density. Shape: (n_spin, n_grid_points)
     """
+    temp = jnp.einsum("sab,raj->srbj", rdm1, grad_ao, precision=precision)
+    return 0.5 * jnp.einsum("srbj,rbj->rs", temp, grad_ao, precision=precision)
 
-    return 0.5 * jnp.einsum("...ab,raj,rbj->r...", rdm1, grad_ao, grad_ao, precision=precision)
-
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def HF_energy_density(
     rdm1: Float[Array, "spin orbitals orbitals"],
     ao: Float[Array, "grid orbitals"],
@@ -540,8 +534,7 @@ def HF_energy_density(
     )
     return vmap(_hf_energy, in_axes=(0, None, 0), out_axes=2)(chi, rdm1, ao)
 
-@jaxtyped
-@typechecked
+@jaxtyped(typechecker=typechecked)
 def HF_density_grad_2_Fock(
     grid: Grid,
     functional: nn.Module,
@@ -612,8 +605,7 @@ def chunked_jvp(chi_tensor, gr_tensor, ao_tensor):
 
     return (jax.jit(chunked_jvp)(chi.transpose(3, 0, 1, 2), gr, ao)).transpose(1, 2, 3, 0)
 
-@jaxtyped
-@typechecked
+@jaxtyped(typechecker=typechecked)
 def HF_coefficient_input_grad_2_Fock(
     grid: Grid,
     functional: nn.Module,
@@ -691,9 +683,8 @@ def abs_clip(arr, threshold):
 
 ######################################################################
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def nonXC(
     rdm1: Float[Array, "orbitals orbitals"],
     h1e: Float[Array, "orbitals orbitals"],
@@ -732,9 +723,8 @@ def nonXC(
 
     return nuclear_repulsion + h1e_energy + coulomb2e_energy
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def one_body_energy(
     rdm1: Float[Array, "orbitals orbitals"],
     h1e: Float[Array, "orbitals orbitals"],
@@ -757,9 +747,8 @@ def one_body_energy(
     return h1e_energy
 
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def coulomb_energy(
     rdm1: Float[Array, "orbitals orbitals"],
     rep_tensor: Float[Array, "orbitals orbitals orbitals orbitals"],
@@ -782,9 +771,8 @@ def coulomb_energy(
     coulomb2e_energy = jnp.einsum("pq,pq->", rdm1, v_coul, precision=precision) / 2.0
     return coulomb2e_energy
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def coulomb_potential(
     rdm1: Float[Array, "orbitals orbitals"],
     rep_tensor: Float[Array, "orbitals orbitals orbitals orbitals"],
@@ -810,9 +798,8 @@ def coulomb_potential(
     """
     return jnp.einsum("pqrt,rt->pq", rep_tensor, rdm1, precision=precision)
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def make_rdm1(
     mo_coeff: Float[Array, "spin orbitals orbitals"],
     mo_occ: Float[Array, "spin orbitals"],
@@ -845,9 +832,8 @@ def make_rdm1(
 
     return jnp.einsum("sij,sj,skj -> sik", mo_coeff, mo_occ, mo_coeff.conj(), precision=precision)
 
-@jaxtyped
-@typechecked
 @jax.jit
+@jaxtyped(typechecker=typechecked)
 def get_occ(
     mo_energies: Float[Array, "spin orbitals"],
     nelecs: Int[Array, "spin"],
diff --git a/grad_dft/popular_functionals.py b/grad_dft/popular_functionals.py
index 69d4c03..e3e7340 100644
--- a/grad_dft/popular_functionals.py
+++ b/grad_dft/popular_functionals.py
@@ -38,7 +38,7 @@ def lsda_x_e(rho: Float[Array, "grid spin"], clip_cte) -> Float[Array, "grid"]:
     -------
     Float[Array, "grid"]
     """
-    rho = jnp.clip(rho, a_min=clip_cte)
+    rho = jnp.clip(rho, clip_cte)
     lda_es = (
         -3.0
         / 4.0
@@ -69,14 +69,14 @@ def b88_x_e(rho: Float[Array, "grid spin"], grad_rho: Float[Array, "grid spin di
 
     beta = 0.0042
 
-    rho = jnp.clip(rho, a_min=clip_cte)
+    rho = jnp.clip(rho, clip_cte)
 
     # LDA preprocessing data: Note that we duplicate the density to sum and divide in the last eq.
-    log_rho = jnp.log2(jnp.clip(rho, a_min=clip_cte))
+    log_rho = jnp.log2(jnp.clip(rho, clip_cte))
 
     grad_rho_norm_sq = jnp.sum(grad_rho**2, axis=-1)
 
-    log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, a_min=clip_cte)) / 2
+    log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, clip_cte)) / 2
 
     # GGA preprocessing data
     log_x_sigma = log_grad_rho_norm - 4 / 3.0 * log_rho
@@ -124,7 +124,7 @@ def pw92_c_e(rho: Float[Array, "grid spin"], clip_cte: float = 1e-30) -> Float[A
     beta3 = jnp.array([[1.6382, 3.3662]])
     beta4 = jnp.array([[0.49294, 0.62517]])
 
-    log_rho = jnp.log2(jnp.clip(rho.sum(axis=1, keepdims=True), a_min=clip_cte))
+    log_rho = jnp.log2(jnp.clip(rho.sum(axis=1, keepdims=True), clip_cte))
     log_rs = jnp.log2((3 / (4 * jnp.pi)) ** (1 / 3)) - log_rho / 3.0
     brs_1_2 = 2 ** (log_rs / 2 + jnp.log2(beta1))
     ars = 2 ** (log_rs + jnp.log2(alpha1))
@@ -161,7 +161,7 @@ def vwn_c_e(rho: Float[Array, "grid spin"], clip_cte: float = 1e-30) -> Float[Ar
     x0 = jnp.array([[-0.10498, -0.325]])
 
     rho = jnp.where(rho > clip_cte, rho, 0.0)
-    log_rho = jnp.log2(jnp.clip(rho.sum(axis=1, keepdims=True), a_min=clip_cte))
+    log_rho = jnp.log2(jnp.clip(rho.sum(axis=1, keepdims=True), clip_cte))
     # assert not jnp.isnan(log_rho).any() and not jnp.isinf(log_rho).any()
     log_rs = jnp.log2((3 / (4 * jnp.pi)) ** (1 / 3)) - log_rho / 3.0
     log_x = log_rs / 2
@@ -232,7 +232,7 @@ def lyp_c_e(rho: Float[Array, "grid spin"], grad_rho: Float[Array, "grid spin 3"
     d = 0.349
     CF = (3 / 10) * (3 * jnp.pi**2) ** (2 / 3)
 
-    rho = jnp.clip(rho, a_min=clip_cte)
+    rho = jnp.clip(rho, clip_cte)
 
     grad_rho_norm_sq = jnp.sum(grad_rho**2, axis=-1)
 
@@ -325,8 +325,7 @@ def b3lyp_exhf_densities(molecule: Molecule, clip_cte: float = 1e-30, *_, **__)
 
     return jnp.stack((lda_e, b88_e, vwn_e, lyp_e), axis=1)
 
-@jaxtyped
-@typechecked
+@jaxtyped(typechecker=typechecked)
 def b3lyp_combine(features: Float[Array, "grid densities"], ehf: Float[Array, "omega spin grid"]) -> Float[Array, "grid densities+1"]:
     r"""
     Auxiliary function to combine the non Hartree-Fock features of B3LYP functional
diff --git a/grad_dft/solid.py b/grad_dft/solid.py
index f6a3ea2..238f6bb 100644
--- a/grad_dft/solid.py
+++ b/grad_dft/solid.py
@@ -266,9 +266,8 @@ def get_mo_grads(self, *args, **kwargs):
         return orbital_grad(self.mo_coeff, self.mo_occ, self.fock, *args, **kwargs)
     
 
-@jaxtyped
-@typechecked
 @partial(jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def one_body_energy(
     rdm1: Complex[Array, "n_kpt n_orbitals n_orbitals"],
     h1e: Complex[Array, "n_kpt n_orbitals n_orbitals"],
@@ -296,9 +295,8 @@ def one_body_energy(
     h1e_energy = jnp.einsum("k,kij,kji->", weights, rdm1, h1e, precision=precision)
     return h1e_energy.real
 
-@jaxtyped
-@typechecked
 @partial(jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def coulomb_potential(
     rdm1: Complex[Array, "n_kpt n_orbitals n_orbitals"],
     rep_tensor: Complex[Array, "n_kpt n_kpt n_orbitals n_orbitals n_orbitals n_orbitals"],
@@ -327,9 +325,8 @@ def coulomb_potential(
     return v_k
     
 
-@jaxtyped
-@typechecked
 @partial(jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def coulomb_energy(
     rdm1: Complex[Array, "n_kpt n_orbitals n_orbitals"],
     rep_tensor: Complex[Array, "n_kpt n_kpt n_orbitals n_orbitals n_orbitals n_orbitals"],
@@ -361,9 +358,8 @@ def coulomb_energy(
     coulomb_energy = jnp.einsum("k,kij,kji->", weights, rdm1, v_k)/2.0
     return coulomb_energy.real
 
-@jaxtyped
-@typechecked
 @partial(jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def non_xc(
     rdm1: Complex[Array, "n_kpt n_orbitals n_orbitals"],
     h1e: Complex[Array, "n_kpt n_orbitals n_orbitals"],
@@ -407,9 +403,8 @@ def non_xc(
     return nuclear_repulsion + kinetic_and_external + coulomb
 
 
-@jaxtyped
-@typechecked
 @partial(jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def make_rdm1(
     mo_coeff: Complex[Array, "n_spin n_kpt n_orbitals n_orbitals"],
     mo_occ: Float[Array, "n_spin n_kpt n_orbitals"],
@@ -433,8 +428,7 @@ def make_rdm1(
     return jnp.einsum("skij,skj,sklj -> skil", mo_coeff, mo_occ, mo_coeff.conj(), precision=precision)
 
 
-@jaxtyped
-@typechecked
+@jaxtyped(typechecker=typechecked)
 def get_occ(
     mo_energies: Float[Array, "n_spin n_kpt n_orbitals"],
     nelecs: Int[Array, "spin"],
@@ -485,9 +479,8 @@ def assign_values(i, mo_occ):
 1BZ. 
 """
 
-@jaxtyped
-@typechecked
 @partial(jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def density(rdm1: Complex[Array, "n_spin n_kpt n_orbitals n_orbitals"], 
             ao: Complex[Array, "n_kpt n_flat_grid n_orbitals"], 
             weights: Float[Array, "n_kpts_or_n_ir_kpts"],
@@ -517,9 +510,8 @@ def density(rdm1: Complex[Array, "n_spin n_kpt n_orbitals n_orbitals"],
     den = jnp.einsum("k,skab,kra,krb->rs", weights, rdm1, ao, ao, precision=precision).real
     return den
 
-@jaxtyped
-@typechecked
 @partial(jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def grad_density(
     rdm1: Complex[Array, "n_spin n_kpt n_orbitals n_orbitals"], 
     ao: Complex[Array, "n_kpt n_flat_grid n_orbitals"], 
@@ -554,9 +546,8 @@ def grad_density(
 
     return 2 * jnp.einsum("k,...kab,kra,krbj->r...j", weights, rdm1, ao, grad_ao, precision=precision).real
 
-@jaxtyped
-@typechecked
 @partial(jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def lapl_density(
     rdm1: Complex[Array, "n_spin n_kpt n_orbitals n_orbitals"], 
     ao: Complex[Array, "n_kpt n_flat_grid n_orbitals"], 
@@ -594,9 +585,8 @@ def lapl_density(
         "k,...kab,kraj,krbj->r...", weights, rdm1, grad_ao, grad_ao, precision=precision
     ) + 2 * jnp.einsum("k,...kab,kra,krbi->r...", weights, rdm1, ao, grad_2_ao, precision=precision)).real
 
-@jaxtyped
-@typechecked
 @partial(jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def kinetic_density(
     rdm1 : Complex[Array, "n_spin n_kpt n_orbitals n_orbitals"],
     grad_ao: Complex[Array, "n_kpt n_flat_grid n_orbitals 3"],
@@ -628,9 +618,8 @@ def kinetic_density(
 
     return 0.5 * jnp.einsum("k,...kab,kraj,krbj->r...", weights, rdm1, grad_ao, grad_ao, precision=precision).real
 
-@jaxtyped
-@typechecked
 @partial(jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def orbital_grad(
         mo_coeff: Complex[Array, "n_spin n_kpt n_orbitals n_orbitals"],
         mo_occ: Float[Array, "n_spin n_kpt n_orbitals"],
diff --git a/grad_dft/train.py b/grad_dft/train.py
index e81680f..6fbcc2a 100644
--- a/grad_dft/train.py
+++ b/grad_dft/train.py
@@ -420,7 +420,7 @@ def dm21_grad_regularization(molecule: Molecule, F: Float[Array, "spin ao ao"])
 
     prefactors = numerator / safe_denominator
 
-    dE = jnp.clip(0.5 * jnp.sum(prefactors * factors), a_min=-10, a_max=10)
+    dE = jnp.clip(0.5 * jnp.sum(prefactors * factors), min=-10, max=10)
 
     return dE**2
 
@@ -579,10 +579,10 @@ def sq_electron_err_int(
         Scalar: the value epsilon described above
     ----------
     """
-    pred_density = jnp.clip(pred_density, a_min=clip_cte)
-    truth_density = jnp.clip(truth_density, a_min=clip_cte)
-    diff_up = jnp.clip(jnp.clip(pred_density[:, 0] - truth_density[:, 0], a_min=clip_cte) ** 2, a_min=clip_cte)
-    diff_dn = jnp.clip(jnp.clip(pred_density[:, 1] - truth_density[:, 1], a_min=clip_cte) ** 2, a_min=clip_cte)
+    pred_density = jnp.clip(pred_density, min=clip_cte)
+    truth_density = jnp.clip(truth_density, min=clip_cte)
+    diff_up = jnp.clip(jnp.clip(pred_density[:, 0] - truth_density[:, 0], min=clip_cte) ** 2, min=clip_cte)
+    diff_dn = jnp.clip(jnp.clip(pred_density[:, 1] - truth_density[:, 1], min=clip_cte) ** 2, min=clip_cte)
     err_int = jnp.sum(diff_up * atoms.grid.weights) + jnp.sum(diff_dn * atoms.grid.weights)
     return err_int
 
diff --git a/requirements.txt b/requirements.txt
index a8f10c0..f9f63ce 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -9,4 +9,4 @@ typeguard==2.13.3
 typing_extensions>=4.8.0
 jaxtyping
 pytest>=7.4.3
-chex>=0.1.91
+chex>=0.1.90
diff --git a/tests/unit/test_kinetic_density_eq.py b/tests/unit/test_kinetic_density_eq.py
new file mode 100644
index 0000000..6e416a1
--- /dev/null
+++ b/tests/unit/test_kinetic_density_eq.py
@@ -0,0 +1,45 @@
+from functools import partial
+import pytest
+
+from grad_dft.molecule import kinetic_density
+import jax.numpy as jnp
+import jax
+from jax.lax import Precision
+import numpy as np
+
+
+@partial(jax.jit, static_argnames="precision")
+def kinetic_density_old(
+    rdm1: jax.Array,
+    grad_ao: jax.Array,
+    precision: Precision = Precision.HIGHEST
+) -> jax.Array:
+    return 0.5 * jnp.einsum(
+        "...ab,raj,rbj->r...", rdm1, grad_ao, grad_ao, precision=precision
+    )
+
+
+@pytest.mark.parametrize("spin", [1, 2])
+@pytest.mark.parametrize("orbitals", [10, 25])
+@pytest.mark.parametrize("grid", [30, 50])
+@pytest.mark.parametrize("precision", [
+    Precision.HIGHEST,
+    Precision.HIGH,
+    Precision.DEFAULT,
+])
+def test_kinetic_density_equivalence(spin, orbitals, grid, precision):
+    rng = jax.random.PRNGKey(0)
+    key1, key2 = jax.random.split(rng)
+
+    rdm1_real = jax.random.randint(key1, (spin, orbitals, orbitals), minval=-80, maxval=80).astype(jnp.float64)
+    rdm1 = 0.5 * (rdm1_real + rdm1_real.transpose(0, 2, 1))
+    grad_ao = jax.random.randint(key2, (grid, orbitals, 3), minval=-80, maxval=80).astype(jnp.float64)
+    result_old = kinetic_density_old(rdm1, grad_ao, precision)
+    result_new = kinetic_density(rdm1, grad_ao, precision)
+
+    np.testing.assert_allclose(
+        result_old, result_new,
+        rtol=1e-5, atol=1e-5,
+        err_msg=f"Results differ for spin={spin}, orbitals={orbitals}, "
+                f"grid={grid}, precision={precision}"
+    )
\ No newline at end of file