pytensor-ml 0.2.2__tar.gz → 0.3.0__tar.gz

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/.github/workflows/run_tests.yml

@@ -35,15 +35,24 @@ jobs:
             paths: |
               tests/test_transformer.py
               tests/test_layers.py
+          - name: recurrent
+            paths: |
               tests/test_recurrent.py
+              tests/rewriting/
+          - name: convolution
+            paths: |
               tests/test_conv.py
               tests/test_conv_transpose.py
               tests/test_padding.py
               tests/test_pool.py
-              tests/rewriting/
+              tests/test_upsample.py
           - name: optim
             paths: |
               tests/optim/
+          - name: models
+            paths: |
+              tests/test_pretrained.py
+              tests/models/
           - name: rest
             paths: |
               tests/test_serialize.py
@@ -51,7 +60,6 @@ jobs:
               tests/test_loss.py
               tests/test_checkpoint.py
               tests/test_model.py
-              tests/test_pretrained.py
               tests/test_pytensorf.py
               tests/test_state.py
               tests/test_base.py
@@ -148,7 +156,7 @@ jobs:
       run: |
         conda activate pytensor_ml
         if [[ $INSTALL_JAX == "1" ]]; then pip install "jax>=0.8,<0.9.1" jaxlib; fi
-        if [[ $INSTALL_MLX == "1" ]]; then pip install "mlx>=0.30,<0.32"; fi
+        if [[ $INSTALL_MLX == "1" ]]; then pip install "mlx>=0.32.2,<0.33"; fi
         if [[ $INSTALL_TORCH == "1" ]]; then pip install torch --index-url https://download.pytorch.org/whl/cpu; fi
       env:
         INSTALL_JAX: ${{ matrix.install-jax }}

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: pytensor_ml
-Version: 0.2.2
+Version: 0.3.0
 Summary: A(nother) deep learning library, built on top of Pytensor
 Project-URL: Homepage, https://github.com/pymc-devs/pytensor-ml
 Project-URL: Documentation, https://pytensor-ml.readthedocs.io
@@ -25,15 +25,16 @@ Classifier: Programming Language :: Python :: 3.14
 Classifier: Topic :: Scientific/Engineering
 Requires-Python: >=3.12
 Requires-Dist: numpy
-Requires-Dist: pytensor<4.0.0,>=3.2.3
-Requires-Dist: safetensors
+Requires-Dist: pytensor<4.0.0,>=3.3.3
 Provides-Extra: dev
+Requires-Dist: cloudpickle; extra == 'dev'
 Requires-Dist: jax; extra == 'dev'
 Requires-Dist: mypy; extra == 'dev'
 Requires-Dist: pre-commit; extra == 'dev'
 Requires-Dist: pytest; extra == 'dev'
 Requires-Dist: pytest-cov; extra == 'dev'
 Requires-Dist: pyyaml; extra == 'dev'
+Requires-Dist: safetensors; extra == 'dev'
 Requires-Dist: scikit-learn; extra == 'dev'
 Provides-Extra: examples
 Requires-Dist: matplotlib; extra == 'examples'
@@ -61,8 +62,9 @@ any other PyTensor graph — a PyMC model included — as there is nothing else
 pip install pytensor-ml
 ```
 
-The only hard dependencies are `pytensor`, `numpy`, and `safetensors`. A backend beyond the default (`numba`,
-`jax`, `torch`, `mlx`) is installed separately, and only loads when you actually compile against it.
+The only hard dependencies are `pytensor` and `numpy`. Saving and loading weights needs `safetensors`
+(`pip install safetensors`). A backend beyond the default (`numba`, `jax`, `torch`, `mlx`) is installed
+separately, and only loads when you actually compile against it.
 
 ## Quickstart
 
@@ -93,11 +95,11 @@ network = Sequential(
     ReLU(),
     Linear("logits", n_in=128, n_out=10),
 )
-model = Model(X_in, network(X_in)).initialize(seed=0)
+model = Model(network(X_in)).initialize(seed=0)
 
 rule = chain(clip_by_global_norm(1.0), adam(learning_rate=cosine_schedule(1e-3, total_steps=500)))
 loss_fn = CrossEntropy(expect_onehot_labels=True, expect_logits=True, reduction="mean")
-step = model.compile_train(rule, loss_fn, ndim_out=2)
+step = model.compile_train(rule, loss_fn)
 
 loader = DataLoader(X, y_onehot, batch_size=64, random_state=0)
 for _ in range(500):

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/README.md

@@ -19,8 +19,9 @@ any other PyTensor graph — a PyMC model included — as there is nothing else
 pip install pytensor-ml
 ```
 
-The only hard dependencies are `pytensor`, `numpy`, and `safetensors`. A backend beyond the default (`numba`,
-`jax`, `torch`, `mlx`) is installed separately, and only loads when you actually compile against it.
+The only hard dependencies are `pytensor` and `numpy`. Saving and loading weights needs `safetensors`
+(`pip install safetensors`). A backend beyond the default (`numba`, `jax`, `torch`, `mlx`) is installed
+separately, and only loads when you actually compile against it.
 
 ## Quickstart
 
@@ -51,11 +52,11 @@ network = Sequential(
     ReLU(),
     Linear("logits", n_in=128, n_out=10),
 )
-model = Model(X_in, network(X_in)).initialize(seed=0)
+model = Model(network(X_in)).initialize(seed=0)
 
 rule = chain(clip_by_global_norm(1.0), adam(learning_rate=cosine_schedule(1e-3, total_steps=500)))
 loss_fn = CrossEntropy(expect_onehot_labels=True, expect_logits=True, reduction="mean")
-step = model.compile_train(rule, loss_fn, ndim_out=2)
+step = model.compile_train(rule, loss_fn)
 
 loader = DataLoader(X, y_onehot, batch_size=64, random_state=0)
 for _ in range(500):

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/conda_envs/environment-docs.yml

@@ -8,7 +8,7 @@ channels:
 dependencies:
   - python>=3.12
   # Runtime deps: autodoc imports pytensor_ml, so the full runtime stack has to be in scope.
-  - pytensor>=3.3.0,<3.4.0
+  - pytensor>=3.3.3,<3.4.0
   - numpy
   - safetensors
   # The gallery extension renders notebook thumbnails with matplotlib.

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/conda_envs/pytensor_ml-gpu_jax.yml

@@ -6,7 +6,7 @@ channels:
 
 dependencies:
   - python>=3.12
-  - pytensor>=3.2.3,<4.0.0
+  - pytensor>=3.3.3,<4.0.0
   - numpy
   - scikit-learn
 

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/conda_envs/pytensor_ml.yml

@@ -5,12 +5,13 @@ channels:
 
 dependencies:
   - python>=3.12
-  - pytensor>=3.3.0,<3.4.0
+  - pytensor>=3.3.3,<3.4.0
   - numpy
   - safetensors
   - scikit-learn
 
   # For testing
+  - cloudpickle  # tests/optim/test_composition.py ships a configured rule to a worker
   - coveralls
   - diff-cover
   - pytest

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/docs/source/api/activations.rst

@@ -10,6 +10,7 @@ Activations
     ReLU
     LeakyReLU
     GELU
+    QuickGELU
     Swish
     Sigmoid
     SoftPlus

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/docs/source/api/layers.rst

@@ -46,6 +46,8 @@ Convolution and pooling
     MaxPool2D
     AvgPool1D
     AvgPool2D
+    Upsample1D
+    Upsample2D
 
 Padding
 -------
@@ -70,6 +72,7 @@ Normalization and regularization
 
     BatchNorm
     LayerNorm
+    GroupNorm
     Dropout
 
 Recurrent

pytensor_ml-0.3.0/docs/source/api/models.rst

@@ -0,0 +1,34 @@
+Pretrained models
+=================
+
+.. currentmodule:: pytensor_ml.models
+
+Dispatch
+--------
+
+.. autosummary::
+    :toctree: generated/
+
+    architecture_name
+    build_from_config
+    register_builder
+
+Blocks
+------
+
+.. autosummary::
+    :toctree: generated/
+
+    ResnetBlock2D
+    AttentionBlock2D
+
+Binding weights
+---------------
+
+.. autosummary::
+    :toctree: generated/
+
+    KeyMap
+    bind_layer_norm
+    bind_linear
+    channels_last

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/docs/source/api/optim.rst

@@ -26,6 +26,16 @@ Update rules
     rprop
     adagrad
     adadelta
+    lbfgs
+
+Line searches
+-------------
+
+.. autosummary::
+    :toctree: generated/
+
+    zoom_line_search
+    LineSearch
 
 Transforms
 ----------
@@ -73,6 +83,7 @@ eliminated the value's last consumer.
 
     constant_schedule
     linear_schedule
+    linear_onecycle_schedule
     cosine_schedule
     exponential_schedule
     polynomial_schedule
@@ -85,9 +96,7 @@ Building blocks
 .. autosummary::
     :toctree: generated/
 
-    counter
     get_gradients
-    reuses_state
     scalar_state
     state_for
     steps_of
@@ -131,3 +140,11 @@ Low-level update functions
     rprop_updates
     adagrad_updates
     adadelta_updates
+    lbfgs_updates
+
+.. currentmodule:: pytensor_ml.optim.lbfgs
+
+.. autosummary::
+    :toctree: generated/
+
+    LBFGSDirection

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/docs/source/api/state.rst

@@ -33,5 +33,7 @@ Initializers
     XavierUniformInitializer
     OrthogonalInitializer
     UnrecordedInitializer
+    EmptyInitializer
+    initial_values_from
     initialize_params
     fans

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/docs/source/api.rst

@@ -14,5 +14,6 @@ API Reference
     api/optim
     api/state
     api/serialization
+    api/models
     api/pytensorf
     api/util

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/docs/source/dev/style_guide.rst

@@ -180,7 +180,7 @@ Every public entrypoint carries an ``Examples`` section. Three rules:
         from pytensor_ml.optim import adam, chain, clip_by_global_norm, compile_train
 
         X = Input("X", shape=(None, 4))
-        loss, target = supervised_loss(Linear("fc", n_in=4, n_out=1)(X), SquaredError(), ndim_out=2)
+        loss, target = supervised_loss(Linear("fc", n_in=4, n_out=1)(X), SquaredError())
 
         step = compile_train(loss, chain(clip_by_global_norm(1.0), adam(1e-3)))
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/docs/source/index.rst

@@ -41,6 +41,9 @@ Quick example
 
     import numpy as np
 
+    from pytensor import config
+    from sklearn.datasets import load_digits
+
     from pytensor_ml.activations import ReLU
     from pytensor_ml.layers import Input, Linear, Sequential
     from pytensor_ml.loss import CrossEntropy
@@ -48,17 +51,21 @@ Quick example
     from pytensor_ml.optim import adam, chain, clip_by_global_norm, cosine_schedule
     from pytensor_ml.util import DataLoader
 
+    X, y = load_digits(return_X_y=True)
+    X = (X / 16.0).astype(config.floatX)
+    y_onehot = np.eye(10, dtype=config.floatX)[y]
+
     X_in = Input("X_in", shape=(None, 64))
     network = Sequential(
         Linear("fc1", n_in=64, n_out=128),
         ReLU(),
         Linear("logits", n_in=128, n_out=10),
     )
-    model = Model(X_in, network(X_in)).initialize(seed=0)
+    model = Model(network(X_in)).initialize(seed=0)
 
     rule = chain(clip_by_global_norm(1.0), adam(learning_rate=cosine_schedule(1e-3, total_steps=500)))
     loss_fn = CrossEntropy(expect_onehot_labels=True, expect_logits=True, reduction="mean")
-    step = model.compile_train(rule, loss_fn, ndim_out=2)
+    step = model.compile_train(rule, loss_fn)
 
     loader = DataLoader(X, y_onehot, batch_size=64, random_state=0)
     for _ in range(500):

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/docs/source/references.bib

@@ -42,3 +42,32 @@
   booktitle = {International Conference on Artificial Intelligence and Statistics},
   year = {2010},
 }
+
+@book{nocedal2006numerical,
+  title = {Numerical Optimization},
+  author = {Nocedal, Jorge and Wright, Stephen J.},
+  edition = {2},
+  publisher = {Springer},
+  address = {New York},
+  year = {2006},
+}
+
+@article{liu1989lbfgs,
+  title = {On the Limited Memory {BFGS} Method for Large Scale Optimization},
+  author = {Liu, Dong C. and Nocedal, Jorge},
+  journal = {Mathematical Programming},
+  volume = {45},
+  number = {1--3},
+  pages = {503--528},
+  year = {1989},
+}
+
+@article{hager2005new,
+  title = {A New Conjugate Gradient Method with Guaranteed Descent and an Efficient Line Search},
+  author = {Hager, William W. and Zhang, Hongchao},
+  journal = {SIAM Journal on Optimization},
+  volume = {16},
+  number = {1},
+  pages = {170--192},
+  year = {2005},
+}

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/docs/source/sharp_edges.rst

@@ -173,6 +173,48 @@ because an endpoint above the starting rate is a legitimate warmup.
 There is no ``warmup`` helper for the same reason: an endpoint above the start already is
 one, and :func:`~pytensor_ml.optim.join_schedules` composes the phases.
 
+Hyperparameters are keyword-only
+--------------------------------
+
+A layer takes its name first and every hyperparameter by keyword. There are no positional
+hyperparameters anywhere in the library, so the torch and keras spelling raises rather than
+binding a hyperparameter to ``name`` and leaving the real one at its default:
+
+.. code-block:: python
+
+    Dropout(0.1)              # TypeError: Dropout's `name` must be a string ... Dropout(p=0.1)
+    Conv2D("conv", 1, 16, 3)  # TypeError: takes 2 positional arguments but 5 were given
+
+    Dropout("drop", p=0.1)    # torch's nn.Dropout(0.1)
+    Conv2D("conv", in_channels=1, out_channels=16, kernel_size=3)
+
+Only the first can tell you which parameter you meant. Once the name slot holds a string
+the extra arguments are anonymous, and the error is the one Python writes.
+
+Every layer's name is optional and falls back to its class, so ``BatchNorm(n_in=4)`` is named
+``BatchNorm`` and a stack of unnamed layers offers several variables of one name. The
+serialization boundary numbers those apart, keying them ``BatchNorm_1_scale``,
+``BatchNorm_2_scale``, ... **in the order the variables are passed**. Collecting them from the
+same graph gives the same order every time, so an unnamed stack round-trips across processes and
+across rebuilds -- but a checkpoint saved from one ordering will not load under another:
+
+.. code-block:: python
+
+    shared = collect_shared_variables(network(X))
+    save_state(shared, "weights.safetensors")
+    load_state(collect_shared_variables(network(X)), "weights.safetensors")   # same order, loads
+
+Name the layers you care about and the numbering never touches them; a name that appears once is
+its own key. The ordinal numbers the layer, not the variable, so two unnamed
+:class:`~pytensor_ml.layers.Linear` layers give ``Linear_1_W``, ``Linear_1_b``, ``Linear_2_W``,
+``Linear_2_b``. Shared state no layer built -- a step counter, an optimizer's moments -- carries no
+layer to number, so a repeat there is numbered at the end of its own name instead.
+
+The only positional argument a layer takes is its name. The two that wrap other layers take those
+instead, and their name by keyword: :class:`~pytensor_ml.layers.Recurrent` is
+``Recurrent(cell, name=...)`` and :class:`~pytensor_ml.layers.Bidirectional` is
+``Bidirectional(forward, backward, name=...)``.
+
 Convolution inputs are channels-last
 ------------------------------------
 

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/examples/gallery/mnist_feed_forward.ipynb

@@ -132,7 +132,7 @@
     ")\n",
     "\n",
     "y_hat = prediction_network(X_in)\n",
-    "model = Model(X_in, y_hat, compile_kwargs={\"mode\": \"NUMBA\"})\n",
+    "model = Model(y_hat, compile_kwargs={\"mode\": \"NUMBA\"})\n",
     "model"
    ]
   },
@@ -170,7 +170,7 @@
     "# `model.compile_train` is the training companion to `model.predict`: it builds the loss against a target,\n",
     "# differentiates it, folds in the BatchNorm running-statistic updates, and compiles a one-step training\n",
     "# function (reusing the model's NUMBA compile settings).\n",
-    "step = model.compile_train(adagrad(learning_rate=1e-3), loss_fn, ndim_out=2)"
+    "step = model.compile_train(adagrad(learning_rate=1e-3), loss_fn)"
    ]
   },
   {

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/pyproject.toml

@@ -35,9 +35,8 @@ keywords = [
 ]
 
 dependencies = [
-  "pytensor>=3.2.3,<4.0.0",
+  "pytensor>=3.3.3,<4.0.0",
   "numpy",
-  "safetensors"
 ]
 
 [project.urls]
@@ -48,6 +47,8 @@ Tracker = "https://github.com/pymc-devs/pytensor-ml/issues"
 
 [project.optional-dependencies]
 dev = [
+    "safetensors",
+    "cloudpickle",
     "pre-commit",
     "pytest",
     "pytest-cov",
@@ -71,6 +72,12 @@ version-file = "pytensor_ml/_version.py"
 [tool.pytest.ini_options]
 minversion = "6.0"
 xfail_strict=true
+# A checkpoint test reads a published snapshot of several GB from disk, so it runs only when asked
+# for: `pytest -m checkpoint`.
+addopts = "-m 'not checkpoint'"
+markers = [
+    "checkpoint: needs a full published checkpoint on disk, not a fixture",
+]
 filterwarnings = [
     "error",
     'ignore:\nArviZ is undergoing a major refactor.*:FutureWarning',
@@ -156,7 +163,7 @@ platforms = ["osx-arm64", "linux-64", "win-64"]
 # the two lists have to move together.
 [tool.pixi.feature.docs.dependencies]
 python = ">=3.12"
-pytensor = ">=3.3.0,<3.4.0"
+pytensor = ">=3.3.3,<3.4.0"
 numpy = "*"
 safetensors = "*"
 # The gallery extension renders notebook thumbnails with matplotlib.

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/pytensor_ml/__init__.py

@@ -3,6 +3,7 @@ import sys
 
 import pytensor_ml.dispatch  # installs the lazy backend-dispatch registration hook
 import pytensor_ml.rewriting.conv  # registers the convolution rewrites into pytensor's databases
+import pytensor_ml.rewriting.layers  # registers the degenerate-dropout rewrite into pytensor's databases
 
 from pytensor_ml._version import __version__
 from pytensor_ml.checkpoint import load_state, save_state

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/pytensor_ml/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.2.2'
-__version_tuple__ = version_tuple = (0, 2, 2)
+__version__ = version = '0.3.0'
+__version_tuple__ = version_tuple = (0, 3, 0)
 
 __commit_id__ = commit_id = None

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/pytensor_ml/activations.py

@@ -151,7 +151,7 @@ class Tanh(Activation):
 
     def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
         out = pt.tanh(x)
-        out.name = "TanH"
+        out.name = "Tanh"
         return out
 
 
@@ -319,10 +319,36 @@ class Swish(Activation):
     def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
         x = pt.as_tensor(x)
         out = x * pt.sigmoid(_constant_like(self.beta, x) * x)
-        out.name = "Swish"
+        out.name = type(self).__name__
         return out
 
 
+class QuickGELU(Swish):
+    r"""
+    Sigmoid approximation to GELU.
+
+    Compute :math:`\mathrm{QuickGELU}(x) = x \, \sigma(1.702 x)`, which is exactly :class:`Swish`
+    at :math:`\beta = 1.702`. Every OpenAI CLIP checkpoint was trained with it, and loading those
+    weights under :class:`GELU` gives quietly wrong activations rather than an error. Prefer
+    :class:`GELU` for anything new.
+
+    Examples
+    --------
+    Use it where a checkpoint was trained with it, such as a CLIP text encoder's MLP:
+
+    .. code-block:: python
+
+        from pytensor_ml.activations import QuickGELU
+        from pytensor_ml.layers import FeedForward, Input
+
+        X = Input("X", shape=(None, 77, 768))
+        activations = FeedForward("mlp", d_model=768, activation=QuickGELU())(X)
+    """
+
+    def __init__(self):
+        super().__init__(beta=1.702)
+
+
 class Softmax(Activation):
     r"""
     Softmax activation.

{pytensor_ml-0.2.2 → pytensor_ml-0.3.0}/pytensor_ml/base.py

@@ -4,6 +4,7 @@ from typing import Protocol, runtime_checkable
 import pytensor.tensor as pt
 
 from pytensor.compile.builders import SymbolicOp
+from pytensor.graph.basic import Variable
 from pytensor.tensor.variable import TensorVariable
 
 
@@ -16,6 +17,39 @@ def _check_input_rank(X: TensorVariable, name: str, n_spatial: int) -> None:
         )
 
 
+def _resolve_layer_name(
+    name: object, default: str, suggested_hyperparameter: str | None = None
+) -> str:
+    """
+    Return the layer's name, rejecting a hyperparameter passed where the name belongs.
+
+    Parameters
+    ----------
+    name : str or None
+        The name the caller passed, annotated as ``object`` because validating it is the point.
+        None selects ``default``.
+    default : str
+        The layer's own name, used when none is given and quoted in the error.
+    suggested_hyperparameter : str, optional
+        The parameter the error tells the caller to pass by keyword instead. Layers with no
+        hyperparameter a name could displace omit it, and the error only reports the bad name.
+
+    Returns
+    -------
+    resolved : str
+        ``name`` when it is a non-empty string, otherwise ``default``.
+    """
+    if name is not None and not isinstance(name, str):
+        message = f"{default}'s `name` must be a string, but got {type(name).__name__} {name!r}."
+        if suggested_hyperparameter is not None:
+            message += (
+                f" Hyperparameters are keyword-only, so pass it by name: "
+                f"{default}({suggested_hyperparameter}={name!r})."
+            )
+        raise TypeError(message)
+    return name if name else default
+
+
 class Layer(ABC):
     """
     Base class for the objects that build layer graphs. Defined here, not in ``pytensor_ml.layers``, so
@@ -89,4 +123,38 @@ class StatefulOp(Protocol):
         """Map each output index to the index of the input that output updates."""
 
 
-__all__ = ["Layer", "LayerOp", "StatefulOp", "UnaryLayerOp"]
+def update_chain_root(variable: Variable) -> tuple[Variable, int] | None:
+    """
+    Trace a stateful op's update input back to the value it ultimately writes, and count the applications
+    behind it.
+
+    Applying one layer object again feeds it what the previous application produced, so the input is the
+    written value itself only at the head of that chain. The depth orders the chain, and the deepest
+    application is the one holding every earlier contribution.
+
+    Parameters
+    ----------
+    variable : Variable
+        The input a stateful op declares itself to update.
+
+    Returns
+    -------
+    root : Variable or None
+        The value written at the head of the chain, or None if the chain breaks at an op that declares no
+        update for the output it reaches.
+    depth : int
+        How many applications separate ``variable`` from ``root``.
+    """
+    depth = 0
+    while True:
+        node = variable.owner
+        if node is None or not isinstance(node.op, StatefulOp):
+            return variable, depth
+        input_index = node.op.update_map().get(node.outputs.index(variable))
+        if input_index is None:
+            return None
+        variable = node.inputs[input_index]
+        depth += 1
+
+
+__all__ = ["Layer", "LayerOp", "StatefulOp", "UnaryLayerOp", "update_chain_root"]