Fix the optional rng_key default in InitFn

blackjax/base.py

@@ -34,7 +34,7 @@ class InitFn(Protocol):
 
     """
 
-    def __call__(self, position: Position, rng_key: PRNGKey | None) -> State:
+    def __call__(self, position: Position, rng_key: PRNGKey | None = None) -> State:
         """Initialize the algorithm's state.
 
         Parameters

tests/test_api_protocols.py

@@ -26,7 +26,7 @@
 import pytest
 
 import blackjax
-from blackjax.base import SamplingAlgorithm
+from blackjax.base import InitFn, SamplingAlgorithm
 from tests.fixtures import std_normal_logdensity
 
 # ---------------------------------------------------------------------------
@@ -36,6 +36,15 @@
 _POSITION = jnp.ones(_DIM)
 
 
+def test_init_protocol_default_matches_hmc_init():
+    protocol_default = inspect.signature(InitFn.__call__).parameters["rng_key"].default
+    hmc_default = (
+        inspect.signature(_make_algorithm("hmc").init).parameters["rng_key"].default
+    )
+
+    assert protocol_default is hmc_default is None
+
+
 def _make_algorithm(name):
     """Instantiate a top-level ``SamplingAlgorithm`` for each public sampler."""
     cov = jnp.eye(_DIM)