matnets 3.2.1__tar.gz → 3.2.3__tar.gz

{matnets-3.2.1 → matnets-3.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: matnets
-Version: 3.2.1
+Version: 3.2.3
 Summary: A small experimental neural network library where neurons are represented as matrices.
 License-Expression: MIT
 License-File: LICENSE

{matnets-3.2.1 → matnets-3.2.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "matnets"
-version = "3.2.1"
+version = "3.2.3"
 description = "A small experimental neural network library where neurons are represented as matrices."
 readme = "docs/index.md"
 requires-python = ">=3.11"

{matnets-3.2.1 → matnets-3.2.3}/src/matnets/activations/__init__.py

@@ -3,6 +3,7 @@ from matnets.activations.exp import (
     sigmoidd,
     softplus,
     softplusd,
+    sss,
     sst,
     tanh,
     tanhd,
@@ -29,6 +30,7 @@ __all__ = [
     "sigmoidd",
     "softplus",
     "softplusd",
+    "sss",
     "sst",
     "tanh",
     "tanhd",

{matnets-3.2.1 → matnets-3.2.3}/src/matnets/activations/exp.py

@@ -56,3 +56,16 @@ def sst(x: jax.Array) -> jax.Array:
     s = jnp.matmul(t, t)
     scale = float(n) ** -2
     return s * scale
+
+
+def sss(x: jax.Array) -> jax.Array:
+    """Scaled squared sigmoid activation.
+
+    Applies element-wise sigmoid, squares the resulting matrices, and
+    scales them by n^-2 where n is the matrix dimension.
+    """
+    n = x.shape[-1]
+    t = jnn.sigmoid(x)
+    s = jnp.matmul(t, t)
+    scale = float(n) ** -2
+    return s * scale

{matnets-3.2.1 → matnets-3.2.3}/src/matnets/nn/recurrent.py

@@ -30,17 +30,26 @@ def lstm_step(
     params: Mapping[str, MatrixParams],
     carry: tuple[Array, Array],
     x: Array,
+    *,
+    activations: tuple[Callable[[Array], Array], Callable[[Array], Array]] = (
+        jax.nn.sigmoid,
+        jnp.tanh,
+    ),
 ) -> tuple[tuple[Array, Array], Array]:
     """LSTM step using ``i``, ``f``, ``g``, and ``o`` dense gate params."""
 
     h, c = carry
+    gate_act, state_act = activations
     combined = jnp.concatenate([h, x], axis=0)
-    i = dense(params["i"], combined, jax.nn.sigmoid)
-    f = dense(params["f"], combined, jax.nn.sigmoid)
-    g = dense(params["g"], combined, jnp.tanh)
-    o = dense(params["o"], combined, jax.nn.sigmoid)
-    next_c = f * c + i * g
-    next_h = o * jnp.tanh(next_c)
+
+    i = dense(params["i"], combined, gate_act)
+    f = dense(params["f"], combined, gate_act)
+    g = dense(params["g"], combined, state_act)
+    o = dense(params["o"], combined, gate_act)
+
+    next_c = jnp.matmul(f, c) + jnp.matmul(i, g)
+    next_h = jnp.matmul(o, state_act(next_c))
+
     return (next_h, next_c), next_h