imt-ring 1.2.1__py3-none-any.whl

ring/ml/optimizer.py

@@ -0,0 +1,149 @@
+from typing import Any, NamedTuple, Optional
+
+import jax
+from jax import lax
+import jax.numpy as jnp
+from jax.tree_util import tree_map
+import optax
+from optax._src import base
+from optax._src import numerics
+from optax._src.transform import add_noise
+from optax._src.transform import AddNoiseState
+
+
+def make_optimizer(
+    lr: float,
+    n_episodes: int,
+    n_steps_per_episode: int,
+    adap_clip: Optional[float] = 0.1,
+    glob_clip: Optional[float] = 0.2,
+    skip_large_update_max_normsq: float = 5.0,
+    skip_large_update_warmup: int = 300,
+    inner_opt=optax.lamb,
+    cos_decay_twice: bool = False,
+    scale_grads: Optional[float] = None,
+    **inner_opt_kwargs,
+):
+    steps = n_steps_per_episode * n_episodes
+    if cos_decay_twice:
+        half_steps = int(steps / 2)
+        schedule = optax.join_schedules(
+            [
+                optax.cosine_decay_schedule(lr, half_steps, 1e-2),
+                optax.cosine_decay_schedule(lr * 1e-2, half_steps),
+            ],
+            [half_steps],
+        )
+    else:
+        schedule = optax.cosine_decay_schedule(lr, steps, 1e-7)
+
+    optimizer = optax.chain(
+        (
+            optax.scale_by_learning_rate(scale_grads, flip_sign=False)
+            if scale_grads is not None
+            else optax.identity()
+        ),
+        (
+            optax.adaptive_grad_clip(adap_clip)
+            if adap_clip is not None
+            else optax.identity()
+        ),
+        optax.clip_by_global_norm(0.2) if glob_clip is not None else optax.identity(),
+        inner_opt(schedule, **inner_opt_kwargs),
+    )
+    optimizer = skip_large_update(
+        optimizer,
+        skip_large_update_max_normsq,
+        max_consecutive_toolarge=6 * 25,
+        warmup=skip_large_update_warmup,
+    )
+    return optimizer
+
+
+class SkipIfLargeUpdatesState(NamedTuple):
+    toolarge_count: jnp.array
+    count: jnp.array
+    inner_state: Any
+    add_noise_state: AddNoiseState
+
+
+def _condition_not_toolarge(updates: base.Updates, max_norm_sq: float):
+    norm_sq = jnp.sum(
+        jnp.array([jnp.sum(p**2) for p in jax.tree_util.tree_leaves(updates)])
+    )
+    # This will also return False if `norm_sq` is NaN or Inf.
+    return norm_sq < max_norm_sq
+
+
+def skip_large_update(
+    inner: base.GradientTransformation,
+    max_norm_sq: float,
+    max_consecutive_toolarge: int,
+    warmup: int = 0,
+    disturb_if_skip: bool = False,
+    disturb_adaptive: bool = False,
+    eta: float = 0.01,
+    gamma: float = 0.55,
+    seed: int = 0,
+) -> base.GradientTransformation:
+    "Also skips NaNs and Infs."
+    inner = base.with_extra_args_support(inner)
+
+    if disturb_adaptive:
+        raise NotImplementedError
+
+    add_noise_transform = add_noise(eta, gamma, seed)
+
+    def init(params):
+        return SkipIfLargeUpdatesState(
+            toolarge_count=jnp.zeros([], jnp.int32),
+            count=jnp.zeros([], jnp.int32),
+            inner_state=inner.init(params),
+            add_noise_state=add_noise_transform.init(params),
+        )
+
+    def update(updates, state: SkipIfLargeUpdatesState, params=None, **extra_args):
+        inner_state = state.inner_state
+        not_toolarge = _condition_not_toolarge(updates, max_norm_sq)
+        toolarge_count = jnp.where(
+            not_toolarge,
+            jnp.zeros([], jnp.int32),
+            numerics.safe_int32_increment(state.toolarge_count),
+        )
+
+        def do_update(updates):
+            updates, new_inner_state = inner.update(
+                updates, inner_state, params, **extra_args
+            )
+            return updates, new_inner_state, state.add_noise_state
+
+        def reject_update(updates):
+            if disturb_if_skip:
+                updates, new_add_noise_state = add_noise_transform.update(
+                    updates, state.add_noise_state, params
+                )
+            else:
+                updates, new_add_noise_state = (
+                    tree_map(jnp.zeros_like, updates),
+                    state.add_noise_state,
+                )
+            return updates, inner_state, new_add_noise_state
+
+        updates, new_inner_state, new_add_noise_state = lax.cond(
+            jnp.logical_or(
+                jnp.logical_or(not_toolarge, toolarge_count > max_consecutive_toolarge),
+                state.count < warmup,
+            ),
+            do_update,
+            reject_update,
+            updates,
+        )
+
+        return updates, SkipIfLargeUpdatesState(
+            toolarge_count=toolarge_count,
+            count=numerics.safe_int32_increment(state.count),
+            inner_state=new_inner_state,
+            add_noise_state=new_add_noise_state,
+        )
+
+    return base.GradientTransformationExtraArgs(init=init, update=update)

ring/ml/ringnet.py

@@ -0,0 +1,279 @@
+from functools import partial
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Callable, Optional
+
+import haiku as hk
+import jax
+import jax.numpy as jnp
+import tree_utils
+
+from ring.maths import safe_normalize
+from ring.ml import base as ml_base
+from ring.utils import pickle_load
+
+
+def _scan_sys(lam: list[int], f):
+    ys = []
+    for i, p in enumerate(lam):
+        ys.append(f(i, p))
+    return tree_utils.tree_batch(ys, backend="jax")
+
+
+def _make_rnno_cell_apply_fn(
+    lam: list[int],
+    inner_cell,
+    send_msg,
+    send_output,
+    hidden_state_dim,
+    message_dim,
+    output_transform: Callable,
+):
+    N = len(lam)
+    parent_array = jnp.array(lam, dtype=jnp.int32)
+
+    def _rnno_cell_apply_fn(inputs, prev_state):
+        empty_message = jnp.zeros((1, message_dim))
+        mailbox = jnp.repeat(empty_message, N, axis=0)
+
+        # message is sent using the hidden state of the last cell
+        # for LSTM `prev_state` is of shape (2 * hidden_state_dim) du to cell state
+        prev_last_hidden_state = prev_state[:, -1, :hidden_state_dim]
+
+        msg = jnp.concatenate(
+            (jax.vmap(send_msg)(prev_last_hidden_state), empty_message)
+        )
+
+        def accumulate_message(link):
+            return jnp.sum(
+                jnp.where(
+                    jnp.repeat((parent_array == link)[:, None], message_dim, axis=-1),
+                    msg[:-1],
+                    mailbox,
+                ),
+                axis=0,
+            )
+
+        mailbox = jax.vmap(accumulate_message)(jnp.arange(N))
+
+        def cell_input(i: int, p: int):
+            local_input = inputs[i]
+            local_cell_input = tree_utils.batch_concat_acme(
+                (local_input, msg[p], mailbox[i]), num_batch_dims=0
+            )
+            return local_cell_input
+
+        stacked_cell_input = _scan_sys(lam, cell_input)
+
+        def update_state(cell_input, state):
+            cell_output, state = inner_cell(cell_input, state)
+            output = output_transform(send_output(cell_output))
+            return output, state
+
+        y, state = jax.vmap(update_state)(stacked_cell_input, prev_state)
+        return y, state
+
+    return _rnno_cell_apply_fn
+
+
+def make_ring(
+    lam: list[int],
+    hidden_state_dim: int = 400,
+    message_dim: int = 200,
+    celltype: str = "gru",
+    stack_rnn_cells: int = 2,
+    send_message_n_layers: int = 1,
+    link_output_dim: int = 4,
+    link_output_normalize: bool = True,
+    link_output_transform: Optional[Callable] = None,
+    layernorm: bool = True,
+) -> SimpleNamespace:
+
+    if link_output_normalize:
+        assert link_output_transform is None
+        link_output_transform = safe_normalize
+    else:
+        if link_output_transform is None:
+            link_output_transform = lambda x: x
+
+    @hk.without_apply_rng
+    @hk.transform_with_state
+    def forward(X):
+        send_msg = hk.nets.MLP(
+            [hidden_state_dim] * send_message_n_layers + [message_dim]
+        )
+
+        inner_cell = StackedRNNCell(
+            celltype, hidden_state_dim, stack_rnn_cells, layernorm=layernorm
+        )
+        send_output = hk.nets.MLP([hidden_state_dim, link_output_dim])
+        state = hk.get_state(
+            "inner_cell_state",
+            [
+                len(lam),
+                stack_rnn_cells,
+                (hidden_state_dim * 2 if celltype == "lstm" else hidden_state_dim),
+            ],
+            init=jnp.zeros,
+        )
+
+        y, state = hk.dynamic_unroll(
+            _make_rnno_cell_apply_fn(
+                lam=lam,
+                inner_cell=inner_cell,
+                send_msg=send_msg,
+                send_output=send_output,
+                hidden_state_dim=hidden_state_dim,
+                message_dim=message_dim,
+                output_transform=link_output_transform,
+            ),
+            X,
+            state,
+        )
+        hk.set_state("inner_cell_state", state)
+        return y
+
+    return forward
+
+
+class StackedRNNCell(hk.Module):
+    def __init__(
+        self,
+        celltype: str,
+        hidden_state_dim,
+        stacks: int,
+        layernorm: bool = False,
+        name: str | None = None,
+    ):
+        super().__init__(name)
+        cell = {"gru": hk.GRU, "lstm": LSTM}[celltype]
+
+        self.cells = [cell(hidden_state_dim) for _ in range(stacks)]
+        self.layernorm = layernorm
+
+    def __call__(self, x, state):
+        output = x
+        next_state = []
+        for i in range(len(self.cells)):
+            output, next_state_i = self.cells[i](output, state[i])
+            next_state.append(next_state_i)
+
+            if self.layernorm:
+                output = hk.LayerNorm(-1, True, True)(output)
+
+        return output, jnp.stack(next_state)
+
+
+class LSTM(hk.RNNCore):
+    def __init__(self, hidden_size: int, name=None):
+        super().__init__(name=name)
+        self.hidden_size = hidden_size
+
+    def __call__(
+        self,
+        inputs: jax.Array,
+        prev_state: jax.Array,
+    ):
+        if len(inputs.shape) > 2 or not inputs.shape:
+            raise ValueError("LSTM input must be rank-1 or rank-2.")
+        prev_state_h = prev_state[: self.hidden_size]
+        prev_state_c = prev_state[self.hidden_size :]
+        x_and_h = jnp.concatenate([inputs, prev_state_h], axis=-1)
+        gated = hk.Linear(4 * self.hidden_size)(x_and_h)
+        i, g, f, o = jnp.split(gated, indices_or_sections=4, axis=-1)
+        f = jax.nn.sigmoid(f + 1)  # Forget bias, as in sonnet.
+        c = f * prev_state_c + jax.nn.sigmoid(i) * jnp.tanh(g)
+        h = jax.nn.sigmoid(o) * jnp.tanh(c)
+        return h, jnp.concatenate((h, c))
+
+    def initial_state(self, batch_size: int | None):
+        raise NotImplementedError
+
+
+class RING(ml_base.AbstractFilter):
+    def __init__(self, params=None, lam=None, jit: bool = True, name=None, **kwargs):
+        self.forward_lam_factory = partial(make_ring, **kwargs)
+        self.params = self._load_params(params)
+        self.lam = lam
+        self._name = name
+
+        if jit:
+            self.apply = jax.jit(self.apply, static_argnames="lam")
+
+    def apply(self, X, params=None, state=None, y=None, lam=None):
+        if lam is None:
+            assert self.lam is not None
+            lam = self.lam
+
+        return super().apply(X, params, state, y, tuple(lam))
+
+    def init(self, bs: Optional[int] = None, X=None, lam=None, seed: int = 1):
+        assert X is not None, "Providing `X` via in `ringnet.init(X=X)` is required"
+        if bs is not None:
+            assert X.ndim == 4
+
+        if X.ndim == 4:
+            if bs is not None:
+                assert bs == X.shape[0]
+            else:
+                bs = X.shape[0]
+            X = X[0]
+
+        # (T, N, F) -> (1, N, F) for faster .init call
+        X = X[0:1]
+
+        if lam is None:
+            assert self.lam is not None
+            lam = self.lam
+
+        key = jax.random.PRNGKey(seed)
+        params, state = self.forward_lam_factory(lam=lam).init(key, X)
+
+        if bs is not None:
+            state = jax.tree_map(lambda arr: jnp.repeat(arr[None], bs, axis=0), state)
+
+        return params, state
+
+    def _apply_batched(self, X, params, state, y, lam):
+        if (params is None and self.params is None) or state is None:
+            _params, _state = self.init(bs=X.shape[0], X=X, lam=lam)
+
+        if params is None and self.params is None:
+            params = _params
+        elif params is None:
+            params = self.params
+        else:
+            pass
+
+        if state is None:
+            state = _state
+
+        yhat, next_state = jax.vmap(
+            self.forward_lam_factory(lam=lam).apply, in_axes=(None, 0, 0)
+        )(params, state, X)
+
+        return yhat, next_state
+
+    @staticmethod
+    def _load_params(params: str | dict | None | Path):
+        assert isinstance(params, (str, dict, type(None), Path))
+        if isinstance(params, (Path, str)):
+            return pickle_load(params)
+        return params
+
+    def nojit(self) -> "RING":
+        ringnet = RING(params=self.params, lam=self.lam, jit=False)
+        ringnet.forward_lam_factory = self.forward_lam_factory
+        return ringnet
+
+    def _pre_save(self, params=None, lam=None) -> None:
+        if params is not None:
+            self.params = params
+        if lam is not None:
+            self.lam = lam
+
+    @staticmethod
+    def _post_load(ringnet: "RING", jit: bool = True) -> "RING":
+        if jit:
+            ringnet.apply = jax.jit(ringnet.apply, static_argnames="lam")
+        return ringnet