Expose the typing module to the public API

lrax/common/algorithm.py

@@ -24,7 +24,7 @@
 from jaxtyping import PRNGKeyArray, PyTree
 from optax import OptState
 
-from .._custom_types import Metrics, Optimizers
+from ..typing import Metrics, Optimizers
 from .env import AbstractEnv, EnvState
 
 type _AlgState = PyTree

lrax/common/callbacks.py

@@ -24,7 +24,7 @@
 import jax.numpy as jnp
 from jaxtyping import ScalarLike
 
-from .._custom_types import Metrics
+from ..typing import Metrics
 
 
 class StopTraining(Exception):

lrax/common/env.py

@@ -23,7 +23,7 @@
 import equinox as eqx
 from jaxtyping import Array, PRNGKeyArray, PyTree
 
-from .._custom_types import Metrics
+from ..typing import Metrics
 
 
 class EnvState(eqx.Module):

lrax/common/kernels.py

@@ -27,18 +27,20 @@
 from equinox.internal import doc_repr
 from jaxtyping import Array, ArrayLike, Float, Scalar, ScalarLike, Shaped
 
-from .._custom_types import DistanceFn, Kernel, P
+from ..typing import DistanceFn, Kernel, Point
 
 _euclidean = doc_repr(lambda x, y: x - y, "<euclidean distance>")
 
 
-def _squared_radius(x: P, y: P, sigma: ArrayLike, distance: DistanceFn) -> Scalar:
+def _squared_radius(
+    x: Point, y: Point, sigma: ArrayLike, distance: DistanceFn
+) -> Scalar:
     return jnp.sum((distance(x, y) / sigma) ** 2)
 
 
 def _radius(
-    x: P,
-    y: P,
+    x: Point,
+    y: Point,
     sigma: ArrayLike,
     distance: DistanceFn,
     order: ScalarLike = 1.0,
@@ -50,8 +52,8 @@ def _radius(
 
 
 def exponential(
-    x: P,
-    y: P,
+    x: Point,
+    y: Point,
     sigma: ArrayLike = 1.0,
     order: ScalarLike = 1.0,
     *,
@@ -92,8 +94,8 @@ def exponential(
 
 
 def energy(
-    x: P,
-    y: P,
+    x: Point,
+    y: Point,
     sigma: ArrayLike = 1.0,
     order: ScalarLike = 1.0,
     *,
@@ -132,8 +134,8 @@ def energy(
 
 
 def crps(
-    x: P,
-    y: P,
+    x: Point,
+    y: Point,
     sigma: ArrayLike = 1.0,
     *,
     distance: DistanceFn = _euclidean,
@@ -170,8 +172,8 @@ def crps(
 
 
 def gaussian(
-    x: P,
-    y: P,
+    x: Point,
+    y: Point,
     sigma: ArrayLike = 1.0,
     *,
     distance: DistanceFn = _euclidean,
@@ -207,8 +209,8 @@ def gaussian(
 
 
 def rational_quadratic(
-    x: P,
-    y: P,
+    x: Point,
+    y: Point,
     sigma: ArrayLike = 1.0,
     alpha: ScalarLike = 1.0,
     *,
@@ -249,8 +251,8 @@ def rational_quadratic(
 
 
 def inverse_multiquadric(
-    x: P,
-    y: P,
+    x: Point,
+    y: Point,
     sigma: ArrayLike = 1.0,
     beta: ScalarLike = 0.5,
     *,
@@ -289,8 +291,8 @@ def inverse_multiquadric(
 
 
 def matern(
-    x: P,
-    y: P,
+    x: Point,
+    y: Point,
     sigma: ArrayLike = 1.0,
     nu: Literal["1/2", "3/2", "5/2"] = "3/2",
     *,
@@ -348,7 +350,7 @@ def matern(
 
 
 def mixture(
-    kernel: Callable[[P, P, Float[ArrayLike, "?*sigma"]], Scalar],
+    kernel: Callable[[Point, Point, Float[ArrayLike, "?*sigma"]], Scalar],
     sigma: Float[ArrayLike, "scales ?*sigma"],
     weights: Float[ArrayLike, " scales"] | None = None,
 ) -> Kernel:
@@ -379,7 +381,7 @@ def mixture(
             f"with shape {jnp.shape(weights)} and a ladder with {rungs} rungs."
         )
 
-    def mixed(x: P, y: P) -> Scalar:
+    def mixed(x: Point, y: Point) -> Scalar:
         def rung(bandwidths: Float[ArrayLike, "?*sigma"]) -> Scalar:
             return kernel(x, y, bandwidths)
 
@@ -395,8 +397,8 @@ def rung(bandwidths: Float[ArrayLike, "?*sigma"]) -> Scalar:
 
 def gram(
     kernel: Kernel,
-    xs: Shaped[P, "n"],
-    ys: Shaped[P, "m"],
+    xs: Shaped[Point, "n"],
+    ys: Shaped[Point, "m"],
     *,
     batch_size: int | None = None,
 ) -> Float[Array, "n m"]:
@@ -417,7 +419,7 @@ def gram(
     the `j`th member of `ys`.
     """
 
-    def row(x: P) -> Float[Array, " m"]:
+    def row(x: Point) -> Float[Array, " m"]:
         return jax.vmap(kernel, in_axes=(None, 0))(x, ys)
 
     if batch_size is None:

lrax/common/logger.py

@@ -27,7 +27,7 @@
 import jax.tree_util as jtu
 import wandb
 
-from .._custom_types import Metrics
+from ..typing import Metrics
 
 
 class Logger(ABC):

lrax/common/networks.py

@@ -30,6 +30,7 @@
 from equinox.internal import doc_repr
 from jaxtyping import Array, PRNGKeyArray, PyTree
 
+from ..typing import Activation
 from .activations import lipswish
 
 _tanh = doc_repr(jnn.tanh, "<function tanh>")
@@ -47,7 +48,7 @@ class SPD(eqx.Module):
 
     size: int = eqx.field(static=True)
     shape: tuple[int, ...] = eqx.field(static=True)
-    metric: Callable
+    metric: Callable[[Array], Array]
     mlp: eqx.nn.MLP
 
     def __init__(
@@ -56,9 +57,9 @@ def __init__(
         diag_size: int,
         width_size: int,
         depth: int,
-        activation: Callable = _lipswish,
-        final_activation: Callable = _identity,
-        metric: Callable = _softplus,
+        activation: Activation = _lipswish,
+        final_activation: Activation = _identity,
+        metric: Callable[[Array], Array] = _softplus,
         *,
         key: PRNGKeyArray,
     ):
@@ -122,7 +123,7 @@ class StackedSPD(eqx.Module):
     size: int = eqx.field(static=True)
     shape: tuple[int, ...] = eqx.field(static=True)
     diagonal_mask: tuple[bool, ...] = eqx.field(static=True)
-    metric: Callable
+    metric: Callable[[Array], Array]
     mlp: eqx.nn.MLP
 
     def __init__(
@@ -132,9 +133,9 @@ def __init__(
         diag_size: int,
         width_size: int,
         depth: int,
-        activation: Callable = _lipswish,
-        final_activation: Callable = _identity,
-        metric: Callable = _softplus,
+        activation: Activation = _lipswish,
+        final_activation: Activation = _identity,
+        metric: Callable[[Array], Array] = _softplus,
         diagonal_mask: Sequence[bool] | None = None,
         *,
         key: PRNGKeyArray,
@@ -239,8 +240,8 @@ def __init__(
         diag_size: int,
         width_size: int,
         depth: int,
-        activation: Callable = _lipswish,
-        final_activation: Callable = _identity,
+        activation: Activation = _lipswish,
+        final_activation: Activation = _identity,
         *,
         key: PRNGKeyArray,
     ):
@@ -299,8 +300,8 @@ def __init__(
         out_size: int | Literal["scalar"],
         width_size: int,
         depth: int,
-        activation: Callable = _lipswish,
-        final_activation: Callable = _tanh,
+        activation: Activation = _lipswish,
+        final_activation: Activation = _tanh,
         *,
         key: PRNGKeyArray,
     ):

lrax/common/policies.py

@@ -18,7 +18,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from collections.abc import Callable, Sequence
+from collections.abc import Sequence
 
 import equinox as eqx
 import jax.nn as jnn
@@ -27,6 +27,8 @@
 from equinox.internal import doc_repr
 from jaxtyping import Array, PRNGKeyArray
 
+from ..typing import Activation
+
 _identity = doc_repr(lambda x: x, "lambda x: x")
 _elu = doc_repr(jnn.elu, "<function elu>")
 
@@ -43,8 +45,8 @@ def __init__(
         act_size: int,
         width_size: int,
         depth: int,
-        activation: Callable = _elu,
-        final_activation: Callable = _identity,
+        activation: Activation = _elu,
+        final_activation: Activation = _identity,
         log_std_init: float = 0.0,
         *,
         key: PRNGKeyArray,
@@ -156,8 +158,8 @@ def __init__(
         obs_size: int,
         width_size: int,
         depth: int,
-        activation: Callable = _elu,
-        final_activation: Callable = _identity,
+        activation: Activation = _elu,
+        final_activation: Activation = _identity,
         *,
         key: PRNGKeyArray,
     ):
@@ -222,8 +224,8 @@ def __init__(
         act_size: int,
         width_size: int,
         depth: int,
-        activation: Callable = _elu,
-        final_activation: Callable = _identity,
+        activation: Activation = _elu,
+        final_activation: Activation = _identity,
         num_critics: int = 2,
         *,
         key: PRNGKeyArray,

lrax/common/scores.py

@@ -28,11 +28,11 @@
 from jax.flatten_util import ravel_pytree
 from jaxtyping import Array, ArrayLike, Float, Scalar, ScalarLike, Shaped
 
-from .._custom_types import Kernel, P
+from ..typing import Kernel, Point
 from .kernels import gaussian
 
 
-def _get_sample_size(samples: Shaped[P, "samples"]) -> int:
+def _get_sample_size(samples: Shaped[Point, "samples"]) -> int:
     def check(path: jtu.KeyPath, leaf: ArrayLike) -> int:
         shape = jnp.shape(leaf)
 
@@ -72,19 +72,19 @@ def _map[X, Y](f: Callable[[X], Y], xs: X, batch_size: int | None) -> Y:
 
 def _total(
     kernel: Kernel,
-    xs: Shaped[P, "n"],
-    ys: Shaped[P, "m"],
+    xs: Shaped[Point, "n"],
+    ys: Shaped[Point, "m"],
     batch_size: int | None,
 ) -> Scalar:
-    def row(x: P) -> Scalar:
+    def row(x: Point) -> Scalar:
         return jnp.sum(jax.vmap(kernel, in_axes=(None, 0))(x, ys))
 
     return jnp.sum(_map(row, xs, batch_size))
 
 
 def _within(
     kernel: Kernel,
-    xs: Shaped[P, "samples"],
+    xs: Shaped[Point, "samples"],
     unbiased: bool,
     batch_size: int | None,
 ) -> Scalar:
@@ -102,15 +102,15 @@ def _within(
     if not unbiased:
         return total / n**2
 
-    def alone(x: P) -> Scalar:
+    def alone(x: Point) -> Scalar:
         return kernel(x, x)
 
     return (total - jnp.sum(_map(alone, xs, batch_size))) / (n * (n - 1))
 
 
 def kernel(
-    observation: P,
-    forecasts: Shaped[P, "samples"],
+    observation: Point,
+    forecasts: Shaped[Point, "samples"],
     kernel: Kernel = gaussian,
     *,
     unbiased: bool = True,
@@ -137,7 +137,7 @@ def kernel(
     strictly proper.
     """
 
-    def against(x: P) -> Scalar:
+    def against(x: Point) -> Scalar:
         return kernel(x, observation)
 
     spread = _within(kernel, forecasts, unbiased, batch_size)
@@ -147,8 +147,8 @@ def against(x: P) -> Scalar:
 
 
 def mmd(
-    observations: Shaped[P, "n"],
-    forecasts: Shaped[P, "m"],
+    observations: Shaped[Point, "n"],
+    forecasts: Shaped[Point, "m"],
     kernel: Kernel = gaussian,
     *,
     unbiased: bool = True,
@@ -194,8 +194,8 @@ def mmd(
 
 
 def variogram(
-    observation: P,
-    forecasts: Shaped[P, "samples"],
+    observation: Point,
+    forecasts: Shaped[Point, "samples"],
     order: ScalarLike = 0.5,
     weights: Float[Array, "d d"] | None = None,
 ) -> Scalar: