Skip parameters the loss cannot reach

pytensor_ml/optim/base.py

@@ -6,7 +6,8 @@
 import pytensor
 
 from pytensor.compile.sharedvalue import SharedVariable
-from pytensor.gradient import DisconnectedInputError, DisconnectedType, grad
+from pytensor.gradient import DisconnectedInputError, grad
+from pytensor.graph.op import io_connection_pattern
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.sharedvar import TensorSharedVariable
 
@@ -82,13 +83,11 @@ def _unreachable_parameter_names(
     loss: TensorVariable, parameters: Sequence[Parameter]
 ) -> list[str]:
     """Name the parameters the loss carries no gradient signal to, which pytensor's own error omits."""
-    gradients = grad(
-        loss, list(parameters), disconnected_inputs="ignore", return_disconnected="disconnected"
-    )
+    connection_pattern = io_connection_pattern(list(parameters), [loss])
     return [
         parameter.name or str(parameter)
-        for parameter, gradient in zip(parameters, gradients)  # type: ignore[arg-type]
-        if isinstance(gradient.type, DisconnectedType)
+        for parameter, to_the_loss in zip(parameters, connection_pattern)
+        if not any(to_the_loss)
     ]
 
 

pytensor_ml/pytensorf/collect.py

@@ -4,6 +4,7 @@
 from pytensor.gradient import DisconnectedGrad, ZeroGrad
 from pytensor.graph import graph_inputs
 from pytensor.graph.basic import Constant, Variable
+from pytensor.graph.op import io_connection_pattern
 from pytensor.graph.traversal import ancestors
 from pytensor.tensor import TensorVariable
 
@@ -57,10 +58,15 @@ def collect_differentiable_params(
     """
     Collect the trainable parameters an optimizer can take gradients of ``outputs`` with respect to.
 
-    These are the parameters :func:`collect_trainable_params` finds, minus any reachable only through a
-    stop-gradient marker -- :func:`pytensor.gradient.disconnected_grad` or
-    :func:`pytensor.gradient.zero_grad` -- which carry no gradient signal by construction. A parameter
-    reached on both a detached and a live path stays in the set, since the live path still carries gradient.
+    Two things disqualify a parameter, and neither subsumes the other. A stop-gradient marker on the only
+    path to it -- :func:`pytensor.gradient.disconnected_grad` or :func:`pytensor.gradient.zero_grad` --
+    though a parameter reached on both a detached and a live path stays, since the live path still carries
+    gradient. Or no connection to the gradient at all, which no marker expresses: an additive constant
+    survives in ``outputs`` as written and vanishes once they differentiate, the shape a physics-informed
+    loss hits when an output bias is gone by the second derivative.
+
+    Both checks are needed because ``zero_grad`` yields a gradient that is connected and zero, which the
+    connection pattern reports as present.
 
     Parameters
     ----------
@@ -72,12 +78,24 @@ def collect_differentiable_params(
     list of TrainableParameter
         The differentiable parameters, in graph-input order.
     """
+    output_list = as_output_list(outputs)
     stop_gradient_outputs = [
         variable
-        for variable in ancestors(as_output_list(outputs))
+        for variable in ancestors(output_list)
         if variable.owner is not None and isinstance(variable.owner.op, DisconnectedGrad | ZeroGrad)
     ]
-    return _collect_inputs_of_type(outputs, TrainableParameter, blockers=stop_gradient_outputs)
+    undetached = _collect_inputs_of_type(
+        outputs, TrainableParameter, blockers=stop_gradient_outputs
+    )
+    if not undetached:
+        return undetached
+
+    connection_pattern = io_connection_pattern(list(undetached), output_list)
+    return [
+        parameter
+        for parameter, to_outputs in zip(undetached, connection_pattern)
+        if any(to_outputs)
+    ]
 
 
 def collect_non_trainable_params(

tests/optim/test_training.py

@@ -3,11 +3,11 @@
 import pytest
 
 from pytensor import config, shared
-from pytensor.gradient import disconnected_grad, zero_grad
+from pytensor.gradient import DisconnectedInputError, disconnected_grad, grad, zero_grad
 from sklearn.datasets import load_digits, make_regression
 from sklearn.preprocessing import MinMaxScaler, OneHotEncoder, StandardScaler
 
-from pytensor_ml.activations import LeakyReLU
+from pytensor_ml.activations import LeakyReLU, Tanh
 from pytensor_ml.layers import BatchNorm2D, Linear, Sequential
 from pytensor_ml.loss import CrossEntropy, SquaredError, supervised_loss
 from pytensor_ml.optim import adam, adamw, compile_train, cosine_schedule, sgd
@@ -309,6 +309,36 @@ def test_compile_train_leaves_a_zero_grad_parameter_untouched_under_weight_decay
     assert moved == {"live_W": True, "live_b": True, "frozen_W": False, "frozen_b": False}
 
 
+def test_a_parameter_that_differentiates_away_is_skipped_unless_you_name_it():
+    """Which policy you get is decided by who chose the parameter set. Collected for you, a parameter the
+    loss cannot reach is left out, so a physics-informed loss trains without hand-enumerating the rest.
+    Named explicitly, it raises -- you asserted it should train, and it cannot."""
+    x = pt.tensor("x", shape=(None, 1))
+    u = Sequential(Linear("hidden", 1, 4), Tanh(), Linear("out", 4, 1))(x)
+    loss = (grad(grad(u.sum(), x).sum(), x) ** 2).mean()
+    every_parameter = collect_trainable_params(loss)
+    initialize(every_parameter)
+    before = {p.name: p.get_value().copy() for p in every_parameter}
+
+    # Decoupled decay is what distinguishes skipped from handed-a-zero: it moves a parameter whatever the
+    # gradient is, so out_b staying put means it is out of the optimizer's set rather than in it with a zero.
+    step = compile_train(loss, adamw(learning_rate=1e-2, weight_decay=0.1), inputs=[x])
+    batch = np.linspace(-1.0, 1.0, 32).reshape(32, 1).astype(config.floatX)
+    history = [float(step(batch)) for _ in range(30)]
+
+    assert history[-1] < history[0]  # training happened, so the assertions below are not vacuous
+    after = {p.name: p.get_value() for p in every_parameter}
+
+    # Exact equality, not a tolerance: skipped means no update expression touched it at all.
+    np.testing.assert_array_equal(after["out_b"], before["out_b"])
+    assert not np.array_equal(
+        after["out_W"], before["out_W"]
+    )  # its sibling in the same layer moved
+
+    with pytest.raises(DisconnectedInputError, match=r"\['out_b'\]"):
+        compile_train(loss, adamw(learning_rate=1e-2), parameters=every_parameter, inputs=[x])
+
+
 def test_extra_updates_write_state_no_gradient_produces():
     # The DQN shape from the other side: a target network kept as a Polyak average of the online weights.
     # No gradient produces that write, so without extra_updates it cannot ride along in the training step.

tests/test_collect.py

@@ -2,9 +2,10 @@
 import pytensor.tensor as pt
 import pytest
 
-from pytensor.gradient import disconnected_grad, zero_grad
+from pytensor.gradient import DisconnectedType, disconnected_grad, grad, zero_grad
 
-from pytensor_ml.layers import Linear
+from pytensor_ml.activations import Tanh
+from pytensor_ml.layers import Linear, Sequential
 from pytensor_ml.params import NonTrainableParameter, TrainableParameter, step_counter
 from pytensor_ml.pytensorf import (
     collect_clock_updates,
@@ -88,6 +89,30 @@ def test_keeps_a_parameter_that_is_also_reached_on_a_live_path(self):
 
         assert names(collect_differentiable_params(loss)) == {"fc_W", "fc_b"}
 
+    def test_excludes_a_parameter_that_differentiates_away_with_no_marker(self):
+        """The physics-informed shape, and the one a graph walk cannot see: an output bias is additive in the
+        network output, so it is an ancestor of the loss as written and gone once the loss differentiates
+        twice with respect to the input. Nothing marks it -- the ops simply propagate nothing to it."""
+        x = pt.tensor("x", shape=(None, 1))
+        u = Sequential(Linear("hidden", 1, 4), Tanh(), Linear("out", 4, 1))(x)
+        loss = (grad(grad(u.sum(), x).sum(), x) ** 2).mean()
+
+        assert names(collect_trainable_params(loss)) == {"hidden_W", "hidden_b", "out_W", "out_b"}
+        assert names(collect_differentiable_params(loss)) == {"hidden_W", "hidden_b", "out_W"}
+
+    def test_a_zero_grad_parameter_is_excluded_though_its_gradient_is_connected(self):
+        """`zero_grad` yields a gradient that is connected and zero, which no connectivity question reports
+        as missing, so the marker walk is what catches it. Replacing that walk with the connection pattern
+        would put these parameters back in the optimizer's set, where weight decay moves them."""
+        X = pt.tensor("X", shape=(None, 4))
+        live_layer, frozen_layer = Linear("live", 4, 2), Linear("frozen", 4, 2)
+        loss = ((live_layer(X) + zero_grad(frozen_layer(X))) ** 2).sum()
+
+        [gradient] = grad(loss, [frozen_layer.W], disconnected_inputs="ignore")
+        assert not isinstance(gradient.type, DisconnectedType)  # connected, just zero
+
+        assert names(collect_differentiable_params(loss)) == {"live_W", "live_b"}
+
 
 class TestCollectNonTrainableParams:
     def test_simple_network_has_none(self, simple_network):