Callback#

class sdm.models.Callback#

Bases: object

Base class for callbacks attached to a model.

requires_grad: ClassVar[bool] = False#

Whether this callback requires gradient calculation.

on_context_preprocessing_end(model: Module, x: TableTensor, y: TableTensor, x_context_only: TableTensor | None, related_tables: RelatedTables[TableTensor] | None) → tuple[TableTensor, TableTensor, TableTensor | None, RelatedTables[TableTensor] | None]#

Called after context preprocessing.

Parameters:

  • model (Module) – The model receiving the callback.

  • x (TableTensor) – The feature tensor of in-context examples with shape [..., R, D] with R rows and D columns.

  • y (TableTensor) – The targets of in-context examples with shape [..., R, 1].

  • x_context_only (TableTensor | None) – The feature tensor available only in context with shape [..., R, D_only], or None when absent.

  • related_tables (RelatedTables[TableTensor] | None) – Related context for in-context examples.

Return type:

tuple[TableTensor, TableTensor, TableTensor | None, RelatedTables[TableTensor] | None]

on_query_preprocessing_end(model: Module, x: TableTensor, related_tables: RelatedTables[TableTensor] | None) → tuple[TableTensor, RelatedTables[TableTensor] | None]#

Called after query preprocessing.

Parameters:

  • model (Module) – the model receiving the callback.

  • x (TableTensor) – The feature tensor of query examples with shape [..., R, D] with R rows and D columns.

  • related_tables (RelatedTables[TableTensor] | None) – Related context for query examples.

Return type:

tuple[TableTensor, RelatedTables[TableTensor] | None]

on_model_forward_end(model: Module, out: TableTensor) → TableTensor#

Called after the model forward pass.

Parameters:

  • model (Module) – the model receiving the callback.

  • out (TableTensor) – The output produced by the model.

Return type:

TableTensor