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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor_ml
-Version: 0.0.6
+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.6 → 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.6'
-__version_tuple__ = version_tuple = (0, 0, 6)
+__version__ = version = '0.0.7'
+__version_tuple__ = version_tuple = (0, 0, 7)
 
 __commit_id__ = commit_id = None

{pytensor_ml-0.0.6 → 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
 
@@ -41,14 +42,21 @@ def save_state(shared_variables: Sequence[SharedVariable], path: str | Path) ->
         Destination archive, written verbatim.
     """
     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."""
+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 value if value.flags["C_CONTIGUOUS"] else np.array(value, order="C")
+    return array if array.flags["C_CONTIGUOUS"] else np.array(array, order="C")
 
 
 def load_state(
@@ -102,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.6 → pytensor_ml-0.0.7}/pytensor_ml/model.py

@@ -63,6 +63,7 @@ class Model:
         *,
         loss: TensorVariable | None = None,
         inputs: Sequence[Variable] | None = None,
+        extra_outputs: Sequence[Variable] | None = None,
     ) -> Function:
         """
         Compile a one-step training function, either against a supervised target or a prebuilt loss.
@@ -70,7 +71,7 @@ class Model:
         Given ``loss_fn``, builds a target placeholder from the model output with :func:`supervised_loss` and
         the step is called as ``step(X_batch, target_batch)``. Given ``loss`` instead, trains that graph
         directly, which is what an autoencoder or a language-model objective needs -- neither has a target
-        separate from its input. Either way the step returns the loss and applies every update in place.
+        separate from its input. Either way the step applies every update in place.
 
         Parameters
         ----------
@@ -87,6 +88,14 @@ class Model:
         inputs : sequence of Variable, optional
             Data inputs of the step, in call order. Collected from ``loss`` when omitted. Belongs to the
             prebuilt path; the supervised path derives its own.
+        extra_outputs : sequence of Variable, optional
+            Diagnostics to return alongside the loss, as in :func:`~pytensor_ml.optim.compile_train`.
+
+        Returns
+        -------
+        step : Function
+            The compiled one-step training function. Returns the loss alone, or ``(loss, *extra_outputs)``
+            when diagnostics were requested.
         """
         if loss_fn is not None:
             if loss is not None or inputs is not None:
@@ -103,6 +112,7 @@ class Model:
             rule,
             parameters=self.weights,
             inputs=inputs,
+            extra_outputs=extra_outputs,
             compile_kwargs=compile_kwargs or self._compile_kwargs,
         )
 

{pytensor_ml-0.0.6 → pytensor_ml-0.0.7}/pytensor_ml/optim/train.py

@@ -19,6 +19,7 @@ def compile_train(
     *,
     parameters: Sequence[Parameter] | None = None,
     inputs: Sequence[Variable] | None = None,
+    extra_outputs: Sequence[Variable] | None = None,
     compile_kwargs: dict | None = None,
 ) -> Function:
     """
@@ -26,8 +27,7 @@ def compile_train(
 
     Differentiates the loss via ``rule``, applies the resulting updates, folds in any non-trainable state
     updates (such as batch-norm running statistics), and compiles. The parameters and data inputs are
-    collected from ``loss`` unless given explicitly. Returns a plain compiled function that maps a batch of
-    inputs to the loss, applying every update in place.
+    collected from ``loss`` unless given explicitly.
 
     Parameters
     ----------
@@ -38,26 +38,40 @@ def compile_train(
     parameters : sequence of shared tensor variable, optional
         Parameters to optimize. Collected from ``loss`` with :func:`collect_trainable_params` when omitted.
     inputs : sequence of Variable, optional
-        Data inputs of the compiled function, in call order. Collected from ``loss`` with
-        :func:`collect_data_inputs` when omitted; pass them explicitly when call order matters (e.g. features
-        before targets).
+        Data inputs of the compiled function, in call order. Collected from ``loss`` and ``extra_outputs``
+        with :func:`collect_data_inputs` when omitted; pass them explicitly when call order matters (e.g.
+        features before targets).
+    extra_outputs : sequence of Variable, optional
+        Diagnostics to return alongside the loss, such as gradient norms or a batch accuracy. Evaluated in
+        the same pass as the gradients, so they see the pre-update parameter values, and they add no
+        non-trainable state updates of their own. A random node reached only through an extra output does
+        still advance its generator, since :func:`~pytensor_ml.pytensorf.function` threads the next-RNG
+        update for every generator the outputs draw from.
     compile_kwargs : dict, optional
         Extra keyword arguments forwarded to the function compiler.
 
     Returns
     -------
-    Function
-        The compiled one-step training function.
+    step : Function
+        The compiled one-step training function, applying every update in place. Returns the loss alone, or
+        ``(loss, *extra_outputs)`` when diagnostics were requested.
     """
+    extra_outputs = list(extra_outputs or [])
+
     if parameters is None:
         parameters = collect_trainable_params(loss)
     if inputs is None:
-        inputs = collect_data_inputs(loss)
+        inputs = collect_data_inputs([loss, *extra_outputs])
 
     updates: dict[SharedVariable, TensorVariable] = dict(rule(loss, parameters))
+
+    # Assigned per key rather than merged: SupportsKeysAndGetItem is invariant in its key type, so
+    # dict.update rejects the narrower NonTrainableParameter keys.
     for parameter, new_value in collect_non_trainable_updates(loss).items():
         updates[parameter] = new_value
 
     require_unique_state_names(updates)
 
-    return function(list(inputs), loss, updates=updates, **(compile_kwargs or {}))
+    outputs = [loss, *extra_outputs] if extra_outputs else loss
+
+    return function(list(inputs), outputs, updates=updates, **(compile_kwargs or {}))

{pytensor_ml-0.0.6 → pytensor_ml-0.0.7}/tests/optim/test_training.py

@@ -120,6 +120,87 @@ def test_compile_train_infers_parameters_and_inputs():
     assert callable(step)
 
 
+def test_compile_train_returns_extra_outputs_in_order():
+    X = pt.tensor("X", shape=(None, 4))
+    prediction = Linear("output", n_in=4, n_out=2)(X)
+    parameters = collect_trainable_params(prediction)
+    initialize(parameters)
+    loss, target = supervised_loss(prediction, SquaredError(), ndim_out=2)
+    weight_norm = pt.sqrt(sum((parameter**2).sum() for parameter in parameters))
+
+    step = compile_train(
+        loss,
+        sgd(1e-2),
+        parameters=parameters,
+        inputs=[X, target],
+        extra_outputs=[weight_norm, prediction],
+    )
+
+    norm_before_step = np.sqrt(sum((parameter.get_value() ** 2).sum() for parameter in parameters))
+
+    rng = np.random.default_rng(0)
+    step_loss, step_norm, step_prediction = step(
+        rng.normal(size=(8, 4)).astype(config.floatX), np.zeros((8, 2), dtype=config.floatX)
+    )
+
+    assert step_loss.ndim == 0
+    assert step_prediction.shape == (8, 2)
+
+    # Extras are evaluated in the gradient pass, so they see the weights the step started from, not the
+    # updated ones.
+    np.testing.assert_allclose(step_norm, norm_before_step, rtol=1e-5)
+
+
+def test_compile_train_collects_inputs_of_extra_outputs():
+    # An extra output may read data the loss never touches; unless those inputs are collected too,
+    # compilation raises MissingInputError. Called by name so the test says nothing about collection order.
+    X = pt.tensor("X", shape=(None, 4))
+    example_weights = pt.tensor("example_weights", shape=(None,))
+    prediction = Linear("output", n_in=4, n_out=2)(X)
+    initialize(collect_trainable_params(prediction))
+    loss, target = supervised_loss(prediction, SquaredError(), ndim_out=2)
+
+    step = compile_train(loss, sgd(1e-2), extra_outputs=[(example_weights * loss).sum()])
+
+    rng = np.random.default_rng(0)
+    step_loss, weighted_loss = step(
+        X=rng.normal(size=(8, 4)).astype(config.floatX),
+        target=np.zeros((8, 2), dtype=config.floatX),
+        example_weights=np.full(8, 2.0, dtype=config.floatX),
+    )
+
+    np.testing.assert_allclose(weighted_loss, step_loss * 16.0, rtol=1e-5)
+
+
+def test_compile_train_ignores_updates_of_extra_outputs():
+    # Extras are read-only observers: a stateful op reached only through an extra output must not have its
+    # write-back folded into the step, or monitoring would silently mutate training state.
+    X = pt.tensor("X", shape=(None, 4))
+    prediction = Linear("fc", n_in=4, n_out=4)(X)
+    monitor = BatchNorm2D("bn", n_in=4)(prediction)
+    parameters = collect_trainable_params(prediction)
+    initialize(parameters)
+    loss, target = supervised_loss(prediction, SquaredError(), ndim_out=2)
+
+    step = compile_train(
+        loss,
+        sgd(1e-2),
+        parameters=parameters,
+        inputs=[X, target],
+        extra_outputs=[monitor.mean()],
+    )
+
+    running_mean = next(
+        p for p in collect_non_trainable_params(monitor) if "running_mean" in p.name
+    )
+    before = running_mean.get_value().copy()
+
+    rng = np.random.default_rng(0)
+    step(rng.normal(size=(16, 4)).astype(config.floatX), np.zeros((16, 4), dtype=config.floatX))
+
+    np.testing.assert_allclose(running_mean.get_value(), before)
+
+
 def test_state_for_requires_named_parameter():
     # An unnamed parameter would leave every state buffer named by its bare slot, so distinct parameters
     # would silently share state at serialization boundaries.

{pytensor_ml-0.0.6 → pytensor_ml-0.0.7}/tests/test_checkpoint.py

@@ -45,7 +45,9 @@ def test_round_trip_restores_parameters_and_optimizer_state(tmp_path):
 
     for _ in range(5):  # drive params and optimizer state away from the snapshot
         step(*batch)
-    assert any(not np.allclose(v.get_value(), snapshot[v.name]) for v in checkpointed)
+    assert any(
+        not np.allclose(variable.get_value(), snapshot[variable.name]) for variable in checkpointed
+    )
 
     load_state(checkpointed, path)
     for variable in checkpointed:
@@ -54,7 +56,7 @@ def test_round_trip_restores_parameters_and_optimizer_state(tmp_path):
 
 def test_step_counter_round_trips_exactly(tmp_path):
     # The step counter is a rank-0 int64 array; the round-trip must preserve its rank, dtype, and value
-    # (np.ascontiguousarray would silently promote it to shape (1,) on save -- see checkpoint._contiguous).
+    # (np.ascontiguousarray would silently promote it to shape (1,) on save -- see _as_saveable_array).
     _, state, _, _ = build_trained_step()
     counter = next(variable for variable in state if variable.name == "adam/step_count")
     saved = counter.get_value()
@@ -71,6 +73,42 @@ def test_step_counter_round_trips_exactly(tmp_path):
     np.testing.assert_array_equal(restored, saved)
 
 
+def test_non_contiguous_value_round_trips_without_reordering(tmp_path):
+    # safetensors serializes an array's raw buffer: handed an F-ordered value it silently writes the
+    # elements in the wrong order rather than raising, so the save must copy to C order first.
+    variable = shared(np.arange(6, dtype="float64").reshape(2, 3).T, "w")
+    assert not variable.get_value(borrow=True).flags["C_CONTIGUOUS"]
+
+    path = tmp_path / "checkpoint.safetensors"
+    save_state([variable], path)
+    variable.set_value(np.zeros((3, 2)))
+    load_state([variable], path)
+
+    np.testing.assert_array_equal(variable.get_value(), np.arange(6).reshape(2, 3).T)
+
+
+@pytest.mark.parametrize(
+    "mode, backend_module", [("JAX", "jax"), ("MLX", "mlx.core")], ids=["jax", "mlx"]
+)
+def test_round_trip_after_compiled_backend_update(mode, backend_module, tmp_path):
+    # A function compiled for a JIT backend stores that backend's array type in the shared variables it
+    # updates, so both sides of a checkpoint taken after training see no numpy at all.
+    pytest.importorskip(backend_module)
+    weight = shared([1.0, 2.0], "w", dtype="float32")
+    counter = shared(0, "step", dtype="int32")
+    pytensor.function([], [], updates={weight: weight * 2, counter: counter + 1}, mode=mode)()
+    assert not isinstance(weight.get_value(), np.ndarray)  # the premise: no longer a numpy value
+
+    path = tmp_path / "checkpoint.safetensors"
+    save_state([weight, counter], path)
+    weight.set_value(np.zeros(2, dtype="float32"))
+    counter.set_value(np.asarray(0, dtype="int32"))
+    load_state([weight, counter], path)
+
+    np.testing.assert_array_equal(weight.get_value(), [2.0, 4.0])
+    np.testing.assert_array_equal(counter.get_value(), 1)
+
+
 def test_name_map_loads_under_renamed_keys(tmp_path):
     saved = shared([1.0, 2.0], "saved/weight")
     path = tmp_path / "checkpoint.safetensors"
@@ -102,7 +140,7 @@ def test_load_rejects_dtype_mismatch(tmp_path):
     # Same shape, different dtype: a real footgun when loading a lower-precision checkpoint into fp64 params.
     save_state([shared([1, 2, 3], "w", dtype="int64")], tmp_path / "c.safetensors")
     target = shared([0.0, 0.0, 0.0], "w")
-    with pytest.raises(ValueError, match="do not match their targets"):
+    with pytest.raises(ValueError, match="archive has int64"):
         load_state([target], tmp_path / "c.safetensors")
     np.testing.assert_array_equal(target.get_value(), [0.0, 0.0, 0.0])
 
@@ -148,4 +186,4 @@ def test_parameters_keep_identity_across_compilation():
     function([X, target], loss, updates=adam(1e-2)(loss, before))
 
     after = collect_trainable_params(prediction)
-    assert all(b is a for b, a in zip(before, after))
+    assert all(b is a for b, a in zip(before, after, strict=True))