pytensor-ml 0.0.5__tar.gz → 0.0.7__tar.gz

{pytensor_ml-0.0.5 → pytensor_ml-0.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor_ml
-Version: 0.0.5
+Version: 0.0.7
 Summary: A(nother) deep learning library, built on top of Pytensor
 Project-URL: Homepage, https://github.com/pymc-devs/pytensor_ml
 Project-URL: Source, https://github.com/pymc-devs/pytensor_ml

{pytensor_ml-0.0.5 → pytensor_ml-0.0.7}/pyproject.toml

@@ -111,6 +111,10 @@ ignore = [
 ]
 fixable = ["ALL"]
 
+[tool.ruff.lint.per-file-ignores]
+# An __init__.py re-exports its package's API and imports submodules for their registration side effects.
+"**/__init__.py" = ["F401"]
+
 
 [tool.ruff.lint.isort]
 lines-between-types = 1

{pytensor_ml-0.0.5 → pytensor_ml-0.0.7}/pytensor_ml/__init__.py

@@ -1,7 +1,7 @@
 import logging
 import sys
 
-import pytensor_ml.dispatch  # noqa: F401  (installs the lazy backend-dispatch registration hook)
+import pytensor_ml.dispatch  # installs the lazy backend-dispatch registration hook
 
 from pytensor_ml._version import __version__
 from pytensor_ml.checkpoint import load_state, save_state

{pytensor_ml-0.0.5 → pytensor_ml-0.0.7}/pytensor_ml/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.0.5'
-__version_tuple__ = version_tuple = (0, 0, 5)
+__version__ = version = '0.0.7'
+__version_tuple__ = version_tuple = (0, 0, 7)
 
 __commit_id__ = commit_id = None

{pytensor_ml-0.0.5 → pytensor_ml-0.0.7}/pytensor_ml/activations.py

@@ -1,7 +1,23 @@
 import numpy as np
 import pytensor.tensor as pt
 
-from pytensor_ml.layers import Layer
+from pytensor import config
+
+from pytensor_ml.base import Layer
+
+
+def _constant_like(value: float, x: pt.TensorVariable) -> pt.TensorVariable:
+    """
+    Wrap a scalar so that combining it with ``x`` cannot widen ``x``'s dtype.
+
+    PyTensor's autocaster types a bare Python float by value, so whether a literal widens its operand
+    depends on that value: against a float32 input ``0.5 * x`` stays float32, while ``0.01 * x``
+    promotes to float64. Pinning the constant to ``x``'s dtype removes the dependence.
+    """
+    dtype = np.dtype(x.dtype)
+    # np.finfo maps complex64 -> float32, keeping complex inputs at their own precision.
+    dtype = np.finfo(dtype).dtype if np.issubdtype(dtype, np.inexact) else np.dtype(config.floatX)
+    return pt.constant(np.asarray(value, dtype=dtype))
 
 
 class Activation(Layer): ...
@@ -46,7 +62,7 @@ class LeakyReLU(Activation):
 
     def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
         x = pt.as_tensor(x)
-        out = pt.switch(x > 0, x, -self.negative_slope * x)
+        out = pt.switch(x > 0, x, _constant_like(self.negative_slope, x) * x)
         out.name = "LeakyReLU"
         return out
 
@@ -133,12 +149,15 @@ class GELU(Activation):
 
     def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
         x = pt.as_tensor(x)
-        # float(...) keeps these Python floats: numpy's float64 scalars would upcast a float32 input to
-        # float64, but a Python float casts by value and preserves the input dtype.
+        half = _constant_like(0.5, x)
+        one = _constant_like(1.0, x)
         if self.approximate:
-            out = 0.5 * x * (1 + pt.tanh(float(np.sqrt(2.0 / np.pi)) * (x + 0.044715 * x**3)))
+            sqrt_2_over_pi = _constant_like(np.sqrt(2.0 / np.pi), x)
+            cubic_coef = _constant_like(0.044715, x)
+            out = half * x * (one + pt.tanh(sqrt_2_over_pi * (x + cubic_coef * x**3)))
         else:
-            out = 0.5 * x * (1 + pt.erf(x / float(np.sqrt(2.0))))
+            sqrt2 = _constant_like(np.sqrt(2.0), x)
+            out = half * x * (one + pt.erf(x / sqrt2))
         out.name = "GELU"
         return out
 
@@ -163,7 +182,7 @@ class Swish(Activation):
 
     def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
         x = pt.as_tensor(x)
-        out = x * pt.sigmoid(self.beta * x)
+        out = x * pt.sigmoid(_constant_like(self.beta, x) * x)
         out.name = "Swish"
         return out
 
@@ -193,4 +212,14 @@ class Softmax(Activation):
         return out
 
 
-__all__ = ["GELU", "LeakyReLU", "ReLU", "Sigmoid", "SoftPlus", "Softmax", "Swish", "Tanh"]
+__all__ = [
+    "GELU",
+    "Activation",
+    "LeakyReLU",
+    "ReLU",
+    "Sigmoid",
+    "SoftPlus",
+    "Softmax",
+    "Swish",
+    "Tanh",
+]

pytensor_ml-0.0.7/pytensor_ml/base.py

@@ -0,0 +1,56 @@
+from abc import ABC, abstractmethod
+from typing import Protocol, runtime_checkable
+
+import pytensor.tensor as pt
+
+from pytensor.compile.builders import SymbolicOp
+from pytensor.tensor.variable import TensorVariable
+
+
+class Layer(ABC):
+    """Base class for the objects that build layer graphs. Defined here, not in ``pytensor_ml.layers``, so
+    that ``pytensor_ml.activations`` can subclass it without a circular import."""
+
+    @abstractmethod
+    def __call__(self, x: pt.TensorLike) -> pt.TensorVariable: ...
+
+
+class LayerOp(SymbolicOp):
+    """Base class for the library's neural-network ops.
+
+    A ``SymbolicOp`` is an ``OpFromGraph`` whose inner graph is rebuilt from its ``__props__`` and input
+    types by :meth:`build_inner_graph`, so equal props with equal inputs yield an identical op. Basing the
+    layers on it (rather than plain ``OpFromGraph``) is what lets the numba backend optimize each inner
+    graph: ``SymbolicOp`` restores the fgraph-aware ``__eq__``/``__hash__`` that a props-carrying
+    ``OpFromGraph`` would otherwise lose, so the ``ofg_inner_graph`` rewrite keeps its optimized inner
+    graph instead of discarding it as unchanged.
+    """
+
+    __props__: tuple[str, ...] = ()
+
+
+class UnaryLayerOp(LayerOp):
+    """A ``LayerOp`` with exactly one output, typed as such.
+
+    ``SymbolicOp.__call__`` is annotated ``Variable | list[Variable]`` because an op may produce many
+    outputs; a unary layer op produces one, so narrow the result to ``TensorVariable``. The ``isinstance``
+    guard narrows for the type checker without a cast and asserts the invariant at runtime.
+    """
+
+    def __call__(self, *inputs, **kwargs) -> TensorVariable:
+        out = super().__call__(*inputs, **kwargs)
+        assert isinstance(out, TensorVariable), f"{type(self).__name__} produced multiple outputs"
+        return out
+
+
+@runtime_checkable
+class StatefulOp(Protocol):
+    """An op that writes some of its outputs back to shared variables it takes as inputs, the way batch
+    norm updates its running statistics. Defining :meth:`update_map` is what marks an op stateful -- the
+    check is structural, so there is nothing to register."""
+
+    def update_map(self) -> dict[int, int]:
+        """Map each output index to the index of the input that output updates."""
+
+
+__all__ = ["Layer", "LayerOp", "StatefulOp", "UnaryLayerOp"]

{pytensor_ml-0.0.5 → pytensor_ml-0.0.7}/pytensor_ml/checkpoint.py

@@ -1,6 +1,7 @@
 from collections.abc import Mapping, Sequence
 from os import fspath
 from pathlib import Path
+from typing import Any
 
 import numpy as np
 
@@ -39,22 +40,23 @@ def save_state(shared_variables: Sequence[SharedVariable], path: str | Path) ->
         Variables whose values to save. Every variable must have a unique, non-``None`` name.
     path : str or pathlib.Path
         Destination archive, written verbatim.
-
-    Raises
-    ------
-    ValueError
-        If any variable is unnamed or two variables share a name.
     """
     indexed = _index_by_name(shared_variables)
-    tensors = {name: _contiguous(variable.get_value()) for name, variable in indexed.items()}
+    tensors = {name: _as_saveable_array(variable.get_value()) for name, variable in indexed.items()}
     save_file(tensors, fspath(path))
 
 
-def _contiguous(value: np.ndarray) -> np.ndarray:
-    """Return a C-contiguous array safetensors will accept, preserving the array's rank."""
-    # np.ascontiguousarray forces ndim >= 1, which would rewrite a rank-0 array (e.g. the step count) to
-    # shape (1,); this keeps rank-0 arrays rank-0.
-    return value if value.flags["C_CONTIGUOUS"] else np.array(value, order="C")
+def _as_saveable_array(value: Any) -> np.ndarray:
+    """
+    Return a value as a C-contiguous numpy array safetensors will accept, preserving its rank.
+
+    A JIT backend stores its own array type in the shared variables a compiled function updates, so a
+    trained model's values are ``jax.Array`` or ``mlx.core.array``; both convert through the array
+    protocol.
+    """
+    array = np.asarray(value)
+    # Not np.ascontiguousarray: it forces ndim >= 1, reshaping a rank-0 array such as the step count.
+    return array if array.flags["C_CONTIGUOUS"] else np.array(array, order="C")
 
 
 def load_state(
@@ -80,12 +82,6 @@ def load_state(
         Maps a variable's name to the archive key to read it from, for loading a checkpoint saved under
         different names (such as HuggingFace's). The mapping must be injective. Names absent from the map
         are matched directly.
-
-    Raises
-    ------
-    ValueError
-        If any variable is unnamed, two variables share a name, ``name_map`` is not injective, the archive
-        keys do not match the target names exactly, or a value's shape or dtype differs from its target.
     """
     indexed = _index_by_name(shared_variables)
     name_map = name_map or {}
@@ -114,7 +110,8 @@ def load_state(
     mismatches = []
     for key, variable in target_by_key.items():
         value = values[key]
-        current = variable.get_value(borrow=True)
+        # Through numpy: a backend array's own dtype object never compares equal to a numpy dtype.
+        current = np.asarray(variable.get_value(borrow=True))
         if value.shape != current.shape or value.dtype != current.dtype:
             mismatches.append(
                 f"  {variable.name!r}: target is {current.dtype} of shape {current.shape}, "

{pytensor_ml-0.0.5 → pytensor_ml-0.0.7}/pytensor_ml/dispatch/__init__.py

@@ -3,11 +3,9 @@ import importlib.abc
 import importlib.util
 import sys
 
-# pytensor auto-loads its own backend dispatches because its linker imports them at compile time -- a
-# third-party op can't hook into that, and pytensor exposes no plugin/entry-point system. To get the same
-# "activate when the backend is actually used, never on the main import path" behavior, watch for pytensor
-# loading a backend's dispatch package and register ours right after. Installing the hook imports nothing
-# heavy; jax/mlx are pulled in only when (and if) that backend compiles a graph.
+# Pytensor imports its own backend dispatches from the linker at compile time and offers no plugin hook, so
+# watch for a backend's dispatch package loading and register ours right after. This keeps jax/mlx off the
+# main import path -- they load only when that backend actually compiles a graph.
 _REGISTRATIONS = {
     "pytensor.link.jax.dispatch": "pytensor_ml.dispatch.jax",
     "pytensor.link.mlx.dispatch": "pytensor_ml.dispatch.mlx",

pytensor_ml-0.0.7/pytensor_ml/dispatch/jax/__init__.py

@@ -0,0 +1,3 @@
+# Imported for its side effect: registering pytensor_ml's JAX funcify dispatches. One submodule per
+# marker op that gets a kernel, mirroring the layout under pytensor_ml/layers.
+import pytensor_ml.dispatch.jax.attention

{pytensor_ml-0.0.5 → pytensor_ml-0.0.7}/pytensor_ml/dispatch/jax/attention.py

@@ -3,7 +3,7 @@ import jax.numpy as jnp
 
 from pytensor.link.jax.dispatch import jax_funcify
 
-from pytensor_ml.attention import AttentionLayer
+from pytensor_ml.layers.attention import AttentionLayer
 
 
 @jax_funcify.register(AttentionLayer)

pytensor_ml-0.0.7/pytensor_ml/dispatch/mlx/__init__.py

@@ -0,0 +1,3 @@
+# Imported for its side effect: registering pytensor_ml's MLX funcify dispatches. One submodule per
+# marker op that gets a kernel, mirroring the layout under pytensor_ml/layers.
+import pytensor_ml.dispatch.mlx.attention

{pytensor_ml-0.0.5 → pytensor_ml-0.0.7}/pytensor_ml/dispatch/mlx/attention.py

@@ -4,7 +4,7 @@ import mlx.core as mx
 
 from pytensor.link.mlx.dispatch import mlx_funcify
 
-from pytensor_ml.attention import AttentionLayer
+from pytensor_ml.layers.attention import AttentionLayer
 
 
 @mlx_funcify.register(AttentionLayer)

pytensor_ml-0.0.7/pytensor_ml/json_serialize.py

@@ -0,0 +1,89 @@
+from collections.abc import Sequence
+
+from pytensor.graph.basic import Variable
+
+# Importing the package registers every op-family handler into the codec's dispatch tables. Those
+# handlers depend on serialize.base, never on this module, so the import runs one way.
+import pytensor_ml.serialize  # noqa: F401
+
+from pytensor_ml.serialize.base import (
+    const_from_json,
+    const_to_json,
+    graph_from_json,
+    graph_to_json,
+    op_from_json,
+    op_to_json,
+    prop_from_json,
+    prop_to_json,
+    qualname,
+    register_from_json,
+    register_type,
+    resolve_class,
+    type_from_json,
+    type_to_json,
+)
+
+
+def serialize_graph(inputs: Sequence[Variable], outputs: Sequence[Variable]) -> dict:
+    """
+    Serialize a pytensor graph to a JSON-native dict.
+
+    Parameters
+    ----------
+    inputs : sequence of Variable
+        The graph's inputs, in order. Include every shared variable the graph reads, since the serialized
+        form addresses inputs positionally and carries no values.
+    outputs : sequence of Variable
+        The graph's outputs.
+
+    Returns
+    -------
+    dict
+        A JSON-native description of the graph's structure (no parameter values).
+    """
+    return graph_to_json(list(inputs), list(outputs))
+
+
+def deserialize_graph(
+    graph_dict: dict, inputs: Sequence[Variable] | None = None
+) -> tuple[list[Variable], list[Variable]]:
+    """
+    Rebuild a pytensor graph from :func:`serialize_graph` output.
+
+    Parameters
+    ----------
+    graph_dict : dict
+        A graph description produced by :func:`serialize_graph`.
+    inputs : sequence of Variable, optional
+        Input leaves to build the graph on, in the original order. Supply these to control input identity —
+        for example to reattach named shared variables. Fresh, plain symbolic variables are created when
+        omitted.
+
+    Returns
+    -------
+    inputs : list of Variable
+        The input variables, in the original order.
+    outputs : list of Variable
+        The rebuilt graph outputs.
+    """
+    return graph_from_json(graph_dict, inputs)
+
+
+__all__ = [
+    "const_from_json",
+    "const_to_json",
+    "deserialize_graph",
+    "graph_from_json",
+    "graph_to_json",
+    "op_from_json",
+    "op_to_json",
+    "prop_from_json",
+    "prop_to_json",
+    "qualname",
+    "register_from_json",
+    "register_type",
+    "resolve_class",
+    "serialize_graph",
+    "type_from_json",
+    "type_to_json",
+]

pytensor_ml-0.0.7/pytensor_ml/layers/__init__.py

@@ -0,0 +1,40 @@
+# Import from the submodule, never the package: this file runs first, so the package is still incomplete.
+# The marker ops are re-exported because saved graphs name them as ``pytensor_ml.layers.<Op>``.
+from pytensor_ml.base import Layer
+from pytensor_ml.layers.attention import (
+    AttentionLayer,
+    CausalSelfAttention,
+    MultiheadAttention,
+    scaled_dot_product_attention,
+)
+from pytensor_ml.layers.combinators import Concatenate, Input, Sequential, Squeeze
+from pytensor_ml.layers.dropout import Dropout, DropoutLayer
+from pytensor_ml.layers.embedding import Embedding, EmbeddingLayer
+from pytensor_ml.layers.linear import Linear, LinearLayer
+from pytensor_ml.layers.norm import (
+    BatchNorm2D,
+    BatchNormLayer,
+    LayerNorm,
+    LayerNormLayer,
+    NoRunningStatsBatchNormLayer,
+    PredictionBatchNormLayer,
+)
+from pytensor_ml.layers.transformer import FeedForward, TransformerBlock
+
+__all__ = [
+    "BatchNorm2D",
+    "CausalSelfAttention",
+    "Concatenate",
+    "Dropout",
+    "Embedding",
+    "FeedForward",
+    "Input",
+    "Layer",
+    "LayerNorm",
+    "Linear",
+    "MultiheadAttention",
+    "Sequential",
+    "Squeeze",
+    "TransformerBlock",
+    "scaled_dot_product_attention",
+]

{pytensor_ml-0.0.5/pytensor_ml → pytensor_ml-0.0.7/pytensor_ml/layers}/attention.py

@@ -4,8 +4,8 @@ import pytensor.tensor as pt
 from pytensor import config
 from pytensor.tensor.variable import TensorVariable
 
-from pytensor_ml.layers import Layer, Linear
-from pytensor_ml.pytensorf import UnaryLayerOp
+from pytensor_ml.base import Layer, UnaryLayerOp
+from pytensor_ml.layers.linear import Linear
 
 
 class AttentionLayer(UnaryLayerOp):
@@ -49,9 +49,8 @@ def _sdpa_graph(
 
     if is_causal:
         sq, sk = q.shape[-2], k.shape[-2]
-        # Position i may attend to j <= i, aligned to the bottom-right when sq != sk so a partial query
-        # block (the typical decoding case) sees the whole prefix. Reduces to a plain lower triangle when
-        # sq == sk
+        # Position i may attend to j <= i, aligned bottom-right so a partial query block (the decoding
+        # case) sees the whole prefix. Reduces to a plain lower triangle when sq == sk.
         q_idx = pt.arange(sq)[:, None]
         k_idx = pt.arange(sk)[None, :]
         causal = pt.where(k_idx <= q_idx + (sk - sq), 0.0, -np.inf).astype(config.floatX)

pytensor_ml-0.0.7/pytensor_ml/layers/combinators.py

@@ -0,0 +1,35 @@
+from collections.abc import Callable
+
+import pytensor.tensor as pt
+
+
+def Input(name: str, shape: tuple[int, ...], dtype: str | None = None) -> pt.TensorVariable:
+    """
+    Create a named symbolic input tensor with a fully static shape.
+
+    Parameters
+    ----------
+    name : str
+        Name of the input variable.
+    shape : tuple of int
+        Static size of each dimension. Raise ``ValueError`` if any entry is not an integer.
+    dtype : str or None
+        Data type of the input. Defaults to ``floatX`` when None.
+    """
+    if not all(isinstance(dim, int) for dim in shape):
+        raise ValueError("All dimensions must be integers")
+
+    return pt.tensor(name=name, shape=shape, dtype=dtype)
+
+
+def Sequential(*layers: Callable) -> Callable:
+    def forward(x: pt.TensorLike) -> pt.TensorLike:
+        for layer in layers:
+            x = layer(x)
+        return x
+
+    return forward
+
+
+Squeeze = pt.squeeze
+Concatenate = pt.concatenate

pytensor_ml-0.0.7/pytensor_ml/layers/dropout.py

@@ -0,0 +1,40 @@
+from typing import Any
+
+import numpy as np
+import pytensor.tensor as pt
+import pytensor.tensor.random as ptr
+
+from pytensor import config
+from pytensor.compile.sharedvalue import shared
+
+from pytensor_ml.base import Layer, UnaryLayerOp
+
+
+class DropoutLayer(UnaryLayerOp):
+    __props__ = ("p",)
+
+    def build_inner_graph(self, X, mask):
+        return [pt.where(mask, ift=X / (1 - self.p), iff=0)]
+
+
+class Dropout(Layer):
+    def __init__(self, name: str | None = None, p: float = 0.5, random_state: Any | None = None):
+        if p < 0.0 or p > 1.0:
+            raise ValueError(f"Dropout probability has to be between 0 and 1, but got {p}")
+        self.name = name if name else "Dropout"
+        self.p = p
+        self.rng = shared(np.random.default_rng(random_state))
+
+    def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
+        X = pt.as_tensor(X)
+        p = pt.as_tensor(self.p, dtype=config.floatX)
+        _, mask = ptr.bernoulli(p=1 - p, size=X.shape, rng=self.rng, return_next_rng=True)
+        mask = mask.astype(config.floatX)
+
+        X_masked = DropoutLayer(
+            name=f"{self.name}[p = {self.p}]",
+            p=self.p,
+        )(X, mask)
+        X_masked.name = f"{self.name}_output"
+
+        return X_masked

pytensor_ml-0.0.7/pytensor_ml/layers/embedding.py

@@ -0,0 +1,53 @@
+import numpy as np
+import pytensor.tensor as pt
+
+from pytensor import config
+
+from pytensor_ml.base import Layer, UnaryLayerOp
+from pytensor_ml.params import trainable
+
+
+class EmbeddingLayer(UnaryLayerOp):
+    __props__ = ("n_embeddings", "n_features")
+
+    def build_inner_graph(self, ids, W):
+        return [W[ids]]
+
+
+class Embedding(Layer):
+    r"""
+    Lookup-table embedding.
+
+    Map each integer index to a learned row of the ``(n_embeddings, n_features)`` table,
+    appending a trailing feature axis of size ``n_features`` while preserving the shape of the
+    index input.
+
+    Parameters
+    ----------
+    name : str or None
+        Name prefix for the layer's parameters. Defaults to "Embedding" when None.
+    n_embeddings : int
+        Number of rows in the table -- the number of distinct indices it can map.
+    n_features : int
+        Size of each embedding row.
+    """
+
+    def __init__(self, name: str | None, n_embeddings: int, n_features: int):
+        self.name = name if name else "Embedding"
+        self.n_embeddings = n_embeddings
+        self.n_features = n_features
+
+        W_value = np.zeros((n_embeddings, n_features), dtype=config.floatX)
+        self.W = trainable(W_value, f"{self.name}_W")
+
+    def __call__(self, ids: pt.TensorLike) -> pt.TensorVariable:
+        ids = pt.as_tensor(ids)
+
+        out = EmbeddingLayer(
+            name=self.name,
+            n_embeddings=self.n_embeddings,
+            n_features=self.n_features,
+        )(ids, self.W)
+        out.name = f"{self.name}_output"
+
+        return out

pytensor_ml-0.0.7/pytensor_ml/layers/linear.py

@@ -0,0 +1,58 @@
+import numpy as np
+import pytensor.tensor as pt
+
+from pytensor import config
+
+from pytensor_ml.base import Layer, UnaryLayerOp
+from pytensor_ml.params import trainable
+
+
+def shape_to_str(shape):
+    inner = ",".join([str(st_dim) if st_dim is not None else "?" for st_dim in shape])
+    return f"({inner})"
+
+
+class LinearLayer(UnaryLayerOp):
+    __props__ = ("n_in", "n_out", "bias")
+
+    def build_inner_graph(self, X, W, b=None):
+        res = X @ W
+        if self.bias:
+            res = res + b
+        return [res]
+
+
+class Linear(Layer):
+    def __init__(self, name: str | None, n_in: int, n_out: int, bias: bool = True):
+        self.name = name if name else "Linear"
+        self.n_in = n_in
+        self.n_out = n_out
+        self.bias = bias
+
+        W_value = np.zeros((n_in, n_out), dtype=config.floatX)
+        self.W = trainable(W_value, f"{self.name}_W")
+
+        if self.bias:
+            b_value = np.zeros(n_out, dtype=config.floatX)
+            self.b = trainable(b_value, f"{self.name}_b")
+
+    def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
+        X = pt.as_tensor(X)
+
+        inputs = [X, self.W]
+        if self.bias:
+            inputs.append(self.b)
+
+        input_shape = shape_to_str(X.type.shape)
+        output_shape = shape_to_str((X @ self.W).type.shape)
+
+        ofg = LinearLayer(
+            name=f"{self.name}[{input_shape} -> {output_shape}]",
+            n_in=self.n_in,
+            n_out=self.n_out,
+            bias=self.bias,
+        )
+        out = ofg(*inputs)
+        out.name = f"{self.name}_output"
+
+        return out