pytensor_ml.model.Model.compile_train

Model.compile_train(rule, loss_fn=None, compile_kwargs=None, *, loss=None, inputs=None, extra_outputs=None)

Compile a one-step training function, either against a supervised target or a prebuilt loss.

Given loss_fn, builds a target placeholder from the model output with supervised_loss() and the step is called with one batch per data input the graph reads, then the target – for the usual single-input model, 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 applies every update in place.

Parameters:

ruleTransform

A configured optimizer, e.g. adam(1e-3).

loss_fnLoss, optional

Callable (target, prediction) -> scalar loss. Mutually exclusive with loss.

compile_kwargsdict, optional

Keyword arguments forwarded to the function compiler. Defaults to the model’s own compile kwargs.

lossTensorVariable, optional

A scalar loss graph built over this model’s output. Mutually exclusive with loss_fn.

inputssequence 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_outputssequence of Variable, optional

Diagnostics to return alongside the loss, as in compile_train().

Returns:

stepFunction

The compiled one-step training function. Returns the loss alone, or (loss, *extra_outputs) when diagnostics were requested.