diff --git a/pytensor_ml/optim/base.py b/pytensor_ml/optim/base.py index f6193f3..cebee36 100644 --- a/pytensor_ml/optim/base.py +++ b/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) ] diff --git a/pytensor_ml/pytensorf/collect.py b/pytensor_ml/pytensorf/collect.py index c2d6ba5..a6251e1 100644 --- a/pytensor_ml/pytensorf/collect.py +++ b/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( diff --git a/tests/optim/test_training.py b/tests/optim/test_training.py index 08c73ae..16f5b04 100644 --- a/tests/optim/test_training.py +++ b/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. diff --git a/tests/test_collect.py b/tests/test_collect.py index 07e478d..3939bef 100644 --- a/tests/test_collect.py +++ b/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):