diff --git a/lrax/common/algorithm.py b/lrax/common/algorithm.py index b47a96c..b99aee5 100644 --- a/lrax/common/algorithm.py +++ b/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 diff --git a/lrax/common/callbacks.py b/lrax/common/callbacks.py index 15583c8..1858a53 100644 --- a/lrax/common/callbacks.py +++ b/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): diff --git a/lrax/common/env.py b/lrax/common/env.py index 4019637..97a5555 100644 --- a/lrax/common/env.py +++ b/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): diff --git a/lrax/common/kernels.py b/lrax/common/kernels.py index 7afd940..a4f377c 100644 --- a/lrax/common/kernels.py +++ b/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, "") -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: diff --git a/lrax/common/logger.py b/lrax/common/logger.py index 2b35b44..089614e 100644 --- a/lrax/common/logger.py +++ b/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): diff --git a/lrax/common/networks.py b/lrax/common/networks.py index 25618b6..63a2ede 100644 --- a/lrax/common/networks.py +++ b/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, "") @@ -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, ): diff --git a/lrax/common/policies.py b/lrax/common/policies.py index 6429f74..11a341d 100644 --- a/lrax/common/policies.py +++ b/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, "") @@ -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, diff --git a/lrax/common/scores.py b/lrax/common/scores.py index 7cb1c1c..cadf08f 100644 --- a/lrax/common/scores.py +++ b/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,11 +72,11 @@ 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)) @@ -84,7 +84,7 @@ def row(x: P) -> Scalar: 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: diff --git a/lrax/common/trainer.py b/lrax/common/trainer.py index 740ad14..204a44d 100644 --- a/lrax/common/trainer.py +++ b/lrax/common/trainer.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 from dataclasses import dataclass from pathlib import Path from typing import Any, Literal, assert_never @@ -28,21 +28,17 @@ import jax.random as jr import jax.tree_util as jtu import optax -from jaxtyping import PRNGKeyArray, PyTree, ScalarLike +from jaxtyping import PRNGKeyArray, PyTree from lightning.pytorch import LightningDataModule from tqdm import tqdm -from .._custom_types import Metrics, Optimizer +from ..typing import LossFn, Optimizer from .algorithm import AbstractAlgorithm from .callbacks import Callback, StopTraining from .env import AbstractEnv from .logger import Logger from .policies import ActorCritic -# takes either (model, batch, keys) or, when `filter_spec` is set, (diff_model, -# static_model, batch, keys); returns ((loss, metrics), grads) -type LossFn = Callable[..., tuple[tuple[ScalarLike, Metrics], PyTree]] - @dataclass class Checkpoint: diff --git a/lrax/ppo/ppo.py b/lrax/ppo/ppo.py index 892332d..2245dbe 100644 --- a/lrax/ppo/ppo.py +++ b/lrax/ppo/ppo.py @@ -29,11 +29,11 @@ from jaxtyping import Array, ArrayLike, PRNGKeyArray, PyTree, ScalarLike from optax import OptState -from .._custom_types import Metrics, Optimizer from .._epsilon import EPSILON from ..common.algorithm import AbstractAlgorithm from ..common.env import AbstractEnv, EnvState from ..common.policies import ActorCritic +from ..typing import Metrics, Optimizer type _AlgState = None diff --git a/lrax/sac/policies.py b/lrax/sac/policies.py index 18375b8..6d820cd 100644 --- a/lrax/sac/policies.py +++ b/lrax/sac/policies.py @@ -18,8 +18,6 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN # THE SOFTWARE. -from collections.abc import Callable - import equinox as eqx import jax.nn as jnn import jax.numpy as jnp @@ -28,6 +26,7 @@ from jaxtyping import Array, PRNGKeyArray from ..common.policies import ContinuousCritic +from ..typing import Activation _identity = doc_repr(lambda x: x, "lambda x: x") _elu = doc_repr(jnn.elu, "") @@ -47,8 +46,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_min: float = -20.0, log_std_max: float = 2.0, *, @@ -159,8 +158,8 @@ def __init__( width_size: int, depth: int, num_critics: int = 2, - activation: Callable = _elu, - final_activation: Callable = _identity, + activation: Activation = _elu, + final_activation: Activation = _identity, log_std_min: float = -20.0, log_std_max: float = 2.0, init_log_alpha: float = 0.0, diff --git a/lrax/sac/sac.py b/lrax/sac/sac.py index 684fb00..9eae191 100644 --- a/lrax/sac/sac.py +++ b/lrax/sac/sac.py @@ -29,11 +29,11 @@ from jaxtyping import Array, PRNGKeyArray, PyTree from optax import OptState -from .._custom_types import Metrics, Optimizer from ..common.algorithm import AbstractAlgorithm from ..common.buffers import ReplayBuffer, Transition from ..common.env import AbstractEnv, EnvState from ..common.policies import ContinuousCritic +from ..typing import Metrics, Optimizer from .policies import ActorCritic type _AlgState = tuple[ContinuousCritic, ReplayBuffer] diff --git a/lrax/shac/shac.py b/lrax/shac/shac.py index a201d84..bba4bed 100644 --- a/lrax/shac/shac.py +++ b/lrax/shac/shac.py @@ -31,11 +31,11 @@ from jaxtyping import Array, PRNGKeyArray, PyTree, ScalarLike from optax import OptState -from .._custom_types import Metrics, Optimizer from .._epsilon import EPSILON from ..common.algorithm import AbstractAlgorithm from ..common.env import AbstractEnv, EnvState from ..common.normalizers import RunningMeanStd +from ..typing import Metrics, Optimizer from .policies import ActorCritic, Critic diff --git a/lrax/_custom_types.py b/lrax/typing.py similarity index 76% rename from lrax/_custom_types.py rename to lrax/typing.py index 01855ad..f7d0c0d 100644 --- a/lrax/_custom_types.py +++ b/lrax/typing.py @@ -27,6 +27,12 @@ type Optimizer = optax.GradientTransformation | optax.GradientTransformationExtraArgs type Optimizers = PyTree[Optimizer] -type P = PyTree[Shaped[ArrayLike, "?*s"], "S"] # type: ignore -type DistanceFn = Callable[[P, P], Array] -type Kernel = Callable[[P, P], Scalar] +# takes either (model, batch, keys) or, when `filter_spec` is set, (diff_model, +# static_model, batch, keys); returns ((loss, metrics), grads) +type LossFn = Callable[..., tuple[tuple[ScalarLike, Metrics], PyTree]] + +type Activation = Callable[[Array], Array] + +type Point = PyTree[Shaped[ArrayLike, "?*s"], "S"] # type: ignore +type DistanceFn = Callable[[Point, Point], Array] +type Kernel = Callable[[Point, Point], Scalar]