Bump to Py3.9 (blackjax-devs#554)

.github/workflows/release.yml

@@ -51,7 +51,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v1
       with:
-        python-version: 3.8
+        python-version: 3.9
     - name: Give PyPI some time to update the index
       run: sleep 240
     - name: Attempt install from PyPI

.github/workflows/test.yml

@@ -24,7 +24,7 @@ jobs:
       - style
     strategy:
       matrix:
-        python-version: [ '3.8', '3.10']
+        python-version: [ '3.9', '3.11']
     steps:
     - uses: actions/checkout@v1
     - name: Set up Python ${{ matrix.python-version }}

.readthedocs.yaml

@@ -1,7 +1,7 @@
 version: 2
 
 python:
-  version: "3.8"
+  version: "3.9"
   install:
       - method: pip
         path: .

blackjax/adaptation/mass_matrix.py

@@ -18,7 +18,7 @@
 parameters used in Hamiltonian Monte Carlo.
 
 """
-from typing import Callable, NamedTuple, Tuple
+from typing import Callable, NamedTuple
 
 import jax
 import jax.numpy as jnp
@@ -68,7 +68,7 @@ class MassMatrixAdaptationState(NamedTuple):
 
 def mass_matrix_adaptation(
     is_diagonal_matrix: bool = True,
-) -> Tuple[Callable, Callable, Callable]:
+) -> tuple[Callable, Callable, Callable]:
     """Adapts the values in the mass matrix by computing the covariance
     between parameters.
 
@@ -156,7 +156,7 @@ def final(mm_state: MassMatrixAdaptationState) -> MassMatrixAdaptationState:
     return init, update, final
 
 
-def welford_algorithm(is_diagonal_matrix: bool) -> Tuple[Callable, Callable, Callable]:
+def welford_algorithm(is_diagonal_matrix: bool) -> tuple[Callable, Callable, Callable]:
     r"""Welford's online estimator of covariance.
 
     It is possible to compute the variance of a population of values in an
@@ -231,7 +231,7 @@ def update(
 
     def final(
         wa_state: WelfordAlgorithmState,
-    ) -> Tuple[Array, int, Array]:
+    ) -> tuple[Array, int, Array]:
         mean, m2, sample_size = wa_state
         covariance = m2 / (sample_size - 1)
         return covariance, sample_size, mean

blackjax/adaptation/pathfinder_adaptation.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Implementation of the Pathinder warmup for the HMC family of sampling algorithms."""
-from typing import Callable, NamedTuple, Tuple, Union
+from typing import Callable, NamedTuple, Union
 
 import jax
 import jax.numpy as jnp
@@ -128,7 +128,7 @@ def update(
             new_ss_state, new_step_size, adaptation_state.inverse_mass_matrix
         )
 
-    def final(warmup_state: PathfinderAdaptationState) -> Tuple[float, Array]:
+    def final(warmup_state: PathfinderAdaptationState) -> tuple[float, Array]:
         """Return the final values for the step size and inverse mass matrix."""
         step_size = jnp.exp(warmup_state.ss_state.log_step_size_avg)
         inverse_mass_matrix = warmup_state.inverse_mass_matrix

blackjax/adaptation/step_size.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Step size adaptation"""
-from typing import Callable, NamedTuple, Tuple
+from typing import Callable, NamedTuple
 
 import jax
 import jax.numpy as jnp
@@ -64,7 +64,7 @@ class DualAveragingAdaptationState(NamedTuple):
 
 def dual_averaging_adaptation(
     target: float, t0: int = 10, gamma: float = 0.05, kappa: float = 0.75
-) -> Tuple[Callable, Callable, Callable]:
+) -> tuple[Callable, Callable, Callable]:
     """Tune the step size in order to achieve a desired target acceptance rate.
 
     Let us note :math:`\\epsilon` the current step size, :math:`\\alpha_t` the

blackjax/adaptation/window_adaptation.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Implementation of the Stan warmup for the HMC family of sampling algorithms."""
-from typing import Callable, List, NamedTuple, Tuple, Union
+from typing import Callable, NamedTuple, Union
 
 import jax
 import jax.numpy as jnp
@@ -45,7 +45,7 @@ class WindowAdaptationState(NamedTuple):
 def base(
     is_mass_matrix_diagonal: bool,
     target_acceptance_rate: float = 0.80,
-) -> Tuple[Callable, Callable, Callable]:
+) -> tuple[Callable, Callable, Callable]:
     """Warmup scheme for sampling procedures based on euclidean manifold HMC.
     The schedule and algorithms used match Stan's :cite:p:`stan_hmc_param` as closely as possible.
 
@@ -191,7 +191,7 @@ def slow_final(warmup_state: WindowAdaptationState) -> WindowAdaptationState:
 
     def update(
         adaptation_state: WindowAdaptationState,
-        adaptation_stage: Tuple,
+        adaptation_stage: tuple,
         position: ArrayLikeTree,
         acceptance_rate: float,
     ) -> WindowAdaptationState:
@@ -233,7 +233,7 @@ def update(
 
         return warmup_state
 
-    def final(warmup_state: WindowAdaptationState) -> Tuple[float, Array]:
+    def final(warmup_state: WindowAdaptationState) -> tuple[float, Array]:
         """Return the final values for the step size and mass matrix."""
         step_size = jnp.exp(warmup_state.ss_state.log_step_size_avg)
         inverse_mass_matrix = warmup_state.imm_state.inverse_mass_matrix
@@ -362,7 +362,7 @@ def build_schedule(
     initial_buffer_size: int = 75,
     final_buffer_size: int = 50,
     first_window_size: int = 25,
-) -> List[Tuple[int, bool]]:
+) -> list[tuple[int, bool]]:
     """Return the schedule for Stan's warmup.
 
     The schedule below is intended to be as close as possible to Stan's :cite:p:`stan_hmc_param`.

blackjax/base.py

@@ -10,7 +10,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Callable, NamedTuple, Tuple
+from typing import Callable, NamedTuple
 
 from typing_extensions import Protocol
 
@@ -64,7 +64,7 @@ class UpdateFn(Protocol):
 
     """
 
-    def __call__(self, rng_key: PRNGKey, state: State) -> Tuple[State, Info]:
+    def __call__(self, rng_key: PRNGKey, state: State) -> tuple[State, Info]:
         """Update the current state using the sampling algorithm.
 
         Parameters

blackjax/mcmc/elliptical_slice.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Public API for the Elliptical Slice sampling Kernel"""
-from typing import Callable, NamedTuple, Tuple
+from typing import Callable, NamedTuple
 
 import jax
 import jax.numpy as jnp
@@ -110,7 +110,7 @@ def kernel(
         rng_key: PRNGKey,
         state: EllipSliceState,
         logdensity_fn: Callable,
-    ) -> Tuple[EllipSliceState, EllipSliceInfo]:
+    ) -> tuple[EllipSliceState, EllipSliceInfo]:
         proposal_generator = elliptical_proposal(
             logdensity_fn, momentum_generator, mean
         )
@@ -205,7 +205,7 @@ def elliptical_proposal(
 
     def generate(
         rng_key: PRNGKey, state: EllipSliceState
-    ) -> Tuple[EllipSliceState, EllipSliceInfo]:
+    ) -> tuple[EllipSliceState, EllipSliceInfo]:
         position, logdensity = state
         key_momentum, key_uniform, key_theta = jax.random.split(rng_key, 3)
         # step 1: sample momentum

blackjax/mcmc/ghmc.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Public API for the Generalized (Non-reversible w/ persistent momentum) HMC Kernel"""
-from typing import Callable, NamedTuple, Tuple
+from typing import Callable, NamedTuple
 
 import jax
 import jax.numpy as jnp
@@ -104,7 +104,7 @@ def kernel(
         momentum_inverse_scale: ArrayLikeTree,
         alpha: float,
         delta: float,
-    ) -> Tuple[GHMCState, hmc.HMCInfo]:
+    ) -> tuple[GHMCState, hmc.HMCInfo]:
         """Generate new sample with the Generalized HMC kernel.
 
         Parameters

blackjax/mcmc/hmc.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Public API for the HMC Kernel"""
-from typing import Callable, NamedTuple, Tuple, Union
+from typing import Callable, NamedTuple, Union
 
 import jax
 
@@ -112,7 +112,7 @@ def kernel(
         step_size: float,
         inverse_mass_matrix: Array,
         num_integration_steps: int,
-    ) -> Tuple[HMCState, HMCInfo]:
+    ) -> tuple[HMCState, HMCInfo]:
         """Generate a new sample with the HMC kernel."""
 
         momentum_generator, kinetic_energy_fn, _ = metrics.gaussian_euclidean(
@@ -281,7 +281,7 @@ def hmc_proposal(
 
     def generate(
         rng_key, state: integrators.IntegratorState
-    ) -> Tuple[integrators.IntegratorState, HMCInfo]:
+    ) -> tuple[integrators.IntegratorState, HMCInfo]:
         """Generate a new chain state."""
         end_state = build_trajectory(state, step_size, num_integration_steps)
         end_state = flip_momentum(end_state)

blackjax/mcmc/mala.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Public API for Metropolis Adjusted Langevin kernels."""
 import operator
-from typing import Callable, NamedTuple, Tuple
+from typing import Callable, NamedTuple
 
 import jax
 import jax.numpy as jnp
@@ -96,7 +96,7 @@ def transition_energy(state, new_state, step_size):
 
     def kernel(
         rng_key: PRNGKey, state: MALAState, logdensity_fn: Callable, step_size: float
-    ) -> Tuple[MALAState, MALAInfo]:
+    ) -> tuple[MALAState, MALAInfo]:
         """Generate a new sample with the MALA kernel."""
         grad_fn = jax.value_and_grad(logdensity_fn)
         integrator = diffusions.overdamped_langevin(grad_fn)

blackjax/mcmc/metrics.py

@@ -27,7 +27,7 @@
 We can also generate a relativistic dynamic :cite:p:`lu2017relativistic`.
 
 """
-from typing import Callable, Tuple
+from typing import Callable
 
 import jax.numpy as jnp
 import jax.scipy as jscipy
@@ -43,7 +43,7 @@
 
 def gaussian_euclidean(
     inverse_mass_matrix: Array,
-) -> Tuple[Callable, EuclideanKineticEnergy, Callable]:
+) -> tuple[Callable, EuclideanKineticEnergy, Callable]:
     r"""Hamiltonian dynamic on euclidean manifold with normally-distributed momentum :cite:p:`betancourt2013general`.
 
     The gaussian euclidean metric is a euclidean metric further characterized

blackjax/mcmc/nuts.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Public API for the NUTS Kernel"""
-from typing import Callable, NamedTuple, Tuple
+from typing import Callable, NamedTuple
 
 import jax
 import jax.numpy as jnp
@@ -121,7 +121,7 @@ def kernel(
         logdensity_fn: Callable,
         step_size: float,
         inverse_mass_matrix: Array,
-    ) -> Tuple[hmc.HMCState, NUTSInfo]:
+    ) -> tuple[hmc.HMCState, NUTSInfo]:
         """Generate a new sample with the NUTS kernel."""
 
         (

blackjax/mcmc/periodic_orbital.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Public API for Periodic Orbital Kernel"""
-from typing import Callable, NamedTuple, Tuple
+from typing import Callable, NamedTuple
 
 import jax
 import jax.numpy as jnp
@@ -142,7 +142,7 @@ def kernel(
         step_size: float,
         inverse_mass_matrix: Array,
         period: int,
-    ) -> Tuple[PeriodicOrbitalState, PeriodicOrbitalInfo]:
+    ) -> tuple[PeriodicOrbitalState, PeriodicOrbitalInfo]:
         """Generate a new orbit with the Periodic Orbital kernel.
 
         Choose a step from the orbit with probability proportional to its weights.
@@ -325,7 +325,7 @@ def periodic_orbital_proposal(
 
     def generate(
         direction: int, init_state: integrators.IntegratorState
-    ) -> Tuple[PeriodicOrbitalState, PeriodicOrbitalInfo]:
+    ) -> tuple[PeriodicOrbitalState, PeriodicOrbitalInfo]:
         """Generate orbit by applying bijection forwards and backwards on period.
 
         As described in algorithm 2 of :cite:p:`neklyudov2022orbital`, each iteration of the periodic orbital

blackjax/mcmc/proposal.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Callable, NamedTuple, Tuple
+from typing import Callable, NamedTuple
 
 import jax
 import jax.numpy as jnp
@@ -42,7 +42,7 @@ class Proposal(NamedTuple):
 
 def proposal_generator(
     energy: Callable, divergence_threshold: float
-) -> Tuple[Callable, Callable]:
+) -> tuple[Callable, Callable]:
     """
 
     Parameters
@@ -61,7 +61,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) -> tuple[Proposal, bool]:
         """Generate a new proposal from a trajectory state.
 
         The trajectory state records information about the position in the state
@@ -95,7 +95,7 @@ def proposal_from_energy_diff(
     new_energy: float,
     divergence_threshold: float,
     state: TrajectoryState,
-) -> Tuple[Proposal, bool]:
+) -> tuple[Proposal, bool]:
     """Computes a new proposal from the energy difference between two states.
 
     It also verifies whether this difference is a divergence, if the
@@ -141,7 +141,7 @@ def asymmetric_proposal_generator(
     transition_energy_fn: Callable,
     divergence_threshold: float,
     proposal_factory: Callable = proposal_from_energy_diff,
-) -> Tuple[Callable, Callable]:
+) -> tuple[Callable, Callable]:
     """A proposal generator that takes into account the transition between
     two states to compute a new proposal.
 
@@ -171,7 +171,7 @@ def update(
         initial_state: TrajectoryState,
         state: TrajectoryState,
         **energy_params,
-    ) -> Tuple[Proposal, bool]:
+    ) -> tuple[Proposal, bool]:
         new_energy = transition_energy_fn(initial_state, state, **energy_params)
         prev_energy = transition_energy_fn(state, initial_state, **energy_params)
         return proposal_factory(prev_energy, new_energy, divergence_threshold, state)