code style

blackjax-devs · Oct 30, 2023 · 6c27125 · 6c27125
1 parent 43ed564
commit 6c27125
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/blackjax/adaptation/chees_adaptation.py b/blackjax/adaptation/chees_adaptation.py
@@ -1,6 +1,6 @@
 """Public API for ChEES-HMC"""
 
-from typing import Callable, NamedTuple
+from typing import Callable, NamedTuple, Optional
 
 import jax
 import jax.numpy as jnp
@@ -254,7 +254,7 @@ def chees_adaptation(
     logprob_fn: Callable,
     num_chains: int,
     *,
-    jitter_seed: int = None,
+    jitter_seed: Optional[int] = None,
 ) -> AdaptationAlgorithm:
     """Adapt the step size and trajectory length (number of integration steps / step size)
     parameters of the jittered HMC algorthm.