fix formatting

blackjax/mcmc/proposal.py

@@ -40,9 +40,7 @@ class Proposal(NamedTuple):
     sum_log_p_accept: float
 
 
-def proposal_generator(
-    energy: Callable
-) -> tuple[Callable, Callable]:
+def proposal_generator(energy: Callable) -> tuple[Callable, Callable]:
     """
 
     Parameters
@@ -59,7 +57,7 @@ def proposal_generator(
     def new(state: TrajectoryState) -> Proposal:
         return Proposal(state, energy(state), 0.0, -jnp.inf)
 
-    def update(initial_energy: float, state: TrajectoryState) -> tuple[Proposal, bool]:
+    def update(initial_energy: float, state: TrajectoryState) -> Proposal:
         """Generate a new proposal from a trajectory state.
 
         The trajectory state records information about the position in the state
@@ -81,9 +79,7 @@ def update(initial_energy: float, state: TrajectoryState) -> tuple[Proposal, boo
 
         """
         new_energy = energy(state)
-        return proposal_from_energy_diff(
-            initial_energy, new_energy, state
-        )
+        return proposal_from_energy_diff(initial_energy, new_energy, state)
 
     return new, update
 
@@ -92,7 +88,7 @@ def proposal_from_energy_diff(
     initial_energy: float,
     new_energy: float,
     state: TrajectoryState,
-) -> tuple[Proposal, bool]:
+) -> Proposal:
     """Computes a new proposal from the energy difference between two states.
 
     Parameters
@@ -118,11 +114,11 @@ def proposal_from_energy_diff(
     sum_log_p_accept = jnp.minimum(delta_energy, 0.0)
 
     return Proposal(
-            state,
-            new_energy,
-            weight,
-            sum_log_p_accept,
-        )
+        state,
+        new_energy,
+        weight,
+        sum_log_p_accept,
+    )
 
 
 def asymmetric_proposal_generator(

blackjax/mcmc/random_walk.py

@@ -63,7 +63,6 @@
 from typing import Callable, NamedTuple, Optional
 
 import jax
-import numpy as np
 from jax import numpy as jnp
 
 from blackjax.base import SamplingAlgorithm

blackjax/mcmc/trajectory.py

@@ -159,9 +159,7 @@ def dynamic_progressive_integration(
         which we say a transition is divergent.
 
     """
-    _, generate_proposal = proposal_generator(
-        hmc_energy(kinetic_energy)
-    )
+    _, generate_proposal = proposal_generator(hmc_energy(kinetic_energy))
     sample_proposal = progressive_uniform_sampling
 
     def integrate(
@@ -320,9 +318,7 @@ def dynamic_recursive_integration(
         Bool to indicate whether to perform additional U turn check between two trajectory.
 
     """
-    _, generate_proposal = proposal_generator(
-        hmc_energy(kinetic_energy)
-    )
+    _, generate_proposal = proposal_generator(hmc_energy(kinetic_energy))
     sample_proposal = progressive_uniform_sampling
 
     def buildtree_integrate(

tests/mcmc/test_proposal_without_chex.py

@@ -29,9 +29,7 @@ def proposal_factory(prev_energy, new_energy, new_state):
             assert new_state == 50
             return new_proposal
 
-        _, update = asymmetric_proposal_generator(
-            transition_energy, proposal_factory
-        )
+        _, update = asymmetric_proposal_generator(transition_energy, proposal_factory)
         proposed = update(30, 50)
         assert proposed == new_proposal