Add type hints to _binom and _b_basis_jvp functions

Author: heptaflar

uncertaintyx/b/jax.py

@@ -6,7 +6,7 @@
 from jax.scipy.special import gammaln
 
 
-def _binom(i, k) -> Array:
+def _binom(i: Array, k: int) -> Array:
     """
     Returns the binomial coefficients for the Bernstein basis
     of degree :math:`k`.
@@ -26,7 +26,7 @@ def _binom(i, k) -> Array:
 )
 
 
-def _b_basis_jvp(k, inputs: tuple[Array], perturbations: tuple[Array]):
+def _b_basis_jvp(k: int, inputs: tuple[Array], perturbations: tuple[Array]) -> tuple[Array, Array]:
     r"""
     Custom forward-mode differentiation (JVP) for the Bernstein basis.