da4ml.converter package

Submodules

da4ml.converter.example module

class da4ml.converter.example.ExampleDAISTracer(model: Callable, hwconf: HWConfig, solver_options: solver_options_t | None = None)

Bases: DAISTracerPluginBase

An example DAIS tracer plugin for the ExampleModel. Two methods must be implemented: - get_input_shapes - apply_model

This plugin must be registered as an entry point under the group dais_tracer.plugins. The entry name should be the module name where the model class is defined. In this case, since the target model class ExampleModel is defined in da4ml.converter.example_plugin, the entry point should be registered under the name da4ml. See pyproject.toml for an example.

apply_model(verbose: bool, inputs: tuple[FixedVariableArray, ...]) → tuple[dict[str, FixedVariableArray], list[str]]

Apply the model and return all intermediate traces.

Parameters:

  • model (The model to be traced.)

  • verbose (Whether to print verbose output.)

  • inputs (Optional inputs to the model.)

Returns:

  • A tuple containing

  • - dict[str, FixedVariableArray] (A dictionary of intermediate names -> FixedVariableArray)

  • - list[str] (A list of output names.)

get_input_shapes()

Get the input shapes for the model. Only used if get_input_kifs returns None.

Return type:

A list of input shapes, or None if not applicable. If cannot be determined, return None.

model: ExampleModel
class da4ml.converter.example.ExampleModel(input_shape: tuple[int, ...] | None = None)

Bases: object

A simple example model class for showcasing DAIS tracer plugin usage.

da4ml.converter.example.operation(inp)

An example operation to be traced. One can use numpy-based operations along with DAIS traceable operations provided in da4ml.trace.ops.

da4ml.converter.plugin module

class da4ml.converter.plugin.DAISTracerPluginBase(model: Callable, hwconf: HWConfig, solver_options: solver_options_t | None = None)

Bases: object

Base class for DAIS tracer plugins.

Methods to be implemented by subclasses: - apply_model - get_input_shapes

apply_model(verbose: bool, inputs: tuple[FixedVariableArray, ...]) → tuple[dict[str, FixedVariableArray], list[str]]

Apply the model and return all intermediate traces.

Parameters:

  • model (The model to be traced.)

  • verbose (Whether to print verbose output.)

  • inputs (Optional inputs to the model.)

Returns:

  • A tuple containing

  • - dict[str, FixedVariableArray] (A dictionary of intermediate names -> FixedVariableArray)

  • - list[str] (A list of output names.)

get_input_shapes() → Sequence[tuple[int, ...]] | None

Get the input shapes for the model. Only used if get_input_kifs returns None.

Return type:

A list of input shapes, or None if not applicable. If cannot be determined, return None.

trace(verbose: bool = False, inputs: tuple[FixedVariableArray, ...] | FixedVariableArray | None = None, inputs_kif: tuple[int, int, int] | None = None, dump: bool = False) → dict[str, FixedVariableArray] | tuple[FixedVariableArray, FixedVariableArray]

Trace the model.

Parameters:

  • verbose (Whether to print verbose output.)

  • inputs (Optional inputs to the model.)

  • inputs_kif (Optional input kif values, only used if inputs is None.)

  • dump (Whether to dump all intermediate traces.)

Returns:

  • If dump is True, returns a dictionary of all intermediate names -> FixedVariableArray.

  • If dump is False, returns a list of output FixedVariableArray.

Module contents

da4ml.converter.get_available_plugins() → dict[str, EntryPoint]
da4ml.converter.trace_model(model: Callable, hwconf: HWConfig | tuple[int, int, int] = HWConfig(1, -1, -1), solver_options: solver_options_t | None = None, verbose: bool = False, inputs: tuple[FixedVariableArray, ...] | FixedVariableArray | None = None, inputs_kif: tuple[int, int, int] | None = None, dump: Literal[False] = False, framework: None | str = None) → tuple[FixedVariableArray, FixedVariableArray]
da4ml.converter.trace_model(model: Callable, hwconf: HWConfig | tuple[int, int, int] = HWConfig(1, -1, -1), solver_options: solver_options_t | None = None, verbose: bool = False, inputs: tuple[FixedVariableArray, ...] | FixedVariableArray | None = None, inputs_kif: tuple[int, int, int] | None = None, dump: Literal[True] = False, framework: None | str = None) → dict[str, FixedVariableArray]