core

Classes

`InputArgs`	Container for input arguments to modules.
`Passage`
`PassageInputAdapter`	PassageInputAdapter(adapter_fn: 'Callable[[InputArgs, Optional[str], Optional[nn.Module]], InputArgs]')
`PassageInputOverrides`
`PassageOutput`	PassageOutput(captured_inputs: 'dict[str, InputArgs]', captured_outputs: 'dict[str, Any]', captured_fake_outputs: 'dict[str, Any]', module_output: 'Any')
`PassageOutputAdapter`	PassageOutputAdapter(adapter_fn: 'Callable[[Any, Optional[str], Optional[nn.Module]], Any]')
`PassageOutputOverrides`
`PatchedModule`

Functions

`always_false_predicate`
`always_true_predicate`
`patch_module`

class InputArgs

Bases: object

Container for input arguments to modules.

__init__(*args, **kwargs)

args: list[Any]

drop_args(index=None)

Parameters:: index (int | slice | None)
Return type:: InputArgs

drop_kwargs(keys=None)

Parameters:: keys (Sequence[str] | None)
Return type:: InputArgs

classmethod from_value(v)

kwargs: dict[str, Any]

exception NoActivePassageContextError: Bases: RuntimeError

class Passage

Bases: Module

__init__(module, *, inputs_to_capture=[], outputs_to_capture=[], input_overrides={}, output_overrides={}, outputs_cache={}, capture_fake_outputs_predicate=<function always_false_predicate>, capture_cache_outputs_predicate=<function always_false_predicate>, early_exit=False, name=None)

Parameters:

module (Module)
inputs_to_capture (Iterable[str])
outputs_to_capture (Iterable[str])
input_overrides (Mapping[str, PassageInputAdapter | InputArgs])
output_overrides (Mapping[str, PassageOutputAdapter | Any])
outputs_cache (dict[str, Any])
capture_fake_outputs_predicate (Callable[[str, Module], bool])
capture_cache_outputs_predicate (Callable[[str, Module], bool])
early_exit (bool)
name (str | None)

active_passages_context = <modelopt.torch.puzzletron.sewing_kit.utils.ActivityContext object>

Parameters:: value (T)
Return type:: ContextManager

classmethod create(module, *, inputs_to_capture=[], outputs_to_capture=[], input_overrides={}, output_overrides={}, outputs_cache={}, capture_fake_outputs_predicate=<function always_false_predicate>, capture_cache_outputs_predicate=<function always_false_predicate>, early_exit=False, name=None)

Parameters:

module (Module)
inputs_to_capture (Iterable[str])
outputs_to_capture (Iterable[str])
input_overrides (Mapping[str, PassageInputAdapter | InputArgs])
output_overrides (Mapping[str, PassageOutputAdapter | Any])
outputs_cache (dict[str, Any])
capture_fake_outputs_predicate (Callable[[str, Module], bool])
capture_cache_outputs_predicate (Callable[[str, Module], bool])
early_exit (bool)
name (str | None)

Return type:

Passage

create_fn_context = <modelopt.torch.puzzletron.sewing_kit.utils.ActivityContext object>

Parameters:: value (T)
Return type:: ContextManager

forward(*args, **kwargs)

Return type:: PassageOutput

freeze()

property input_overrides: PassageInputOverrides

is_active()

Return type:: bool

classmethod module_name_relative_to_active_passage(module)

Parameters:: module (PatchedModule)
Return type:: str

property output_overrides: PassageOutputOverrides

reset()

run(*args, **kwargs)

Return type:: PassageOutput

unfreeze()

class PassageInputAdapter

Bases: object

PassageInputAdapter(adapter_fn: ‘Callable[[InputArgs, Optional[str], Optional[nn.Module]], InputArgs]’)

__init__(adapter_fn)

Parameters:: adapter_fn (Callable[[InputArgs, str | None, Module | None], InputArgs])
Return type:: None

adapter_fn: Callable[[InputArgs, str | None, Module | None], InputArgs]

class PassageInputOverrides

Bases: dict[str, PassageInputAdapter | InputArgs]

__init__(input_overrides={})

Parameters:: input_overrides (Mapping[str, PassageInputAdapter | InputArgs])

class PassageOutput

Bases: object

PassageOutput(captured_inputs: ‘dict[str, InputArgs]’, captured_outputs: ‘dict[str, Any]’, captured_fake_outputs: ‘dict[str, Any]’, module_output: ‘Any’)

__init__(captured_inputs, captured_outputs, captured_fake_outputs, module_output)

Parameters:

captured_inputs (dict[str, InputArgs])
captured_outputs (dict[str, Any])
captured_fake_outputs (dict[str, Any])
module_output (Any)

Return type:

None

captured_fake_outputs: dict[str, Any]

captured_inputs: dict[str, InputArgs]

captured_outputs: dict[str, Any]

module_output: Any

class PassageOutputAdapter

Bases: object

PassageOutputAdapter(adapter_fn: ‘Callable[[Any, Optional[str], Optional[nn.Module]], Any]’)

__init__(adapter_fn)

Parameters:: adapter_fn (Callable[[Any, str | None, Module | None], Any])
Return type:: None

adapter_fn: Callable[[Any, str | None, Module | None], Any]

class PassageOutputOverrides

Bases: dict[str, PassageOutputAdapter | Any]

__init__(output_overrides={})

Parameters:: output_overrides (Mapping[str, PassageOutputAdapter | Any])

class PatchedModule: Bases: object

exception RequiredPassageOutputsCapturedSignal: Bases: Exception

always_false_predicate(module_name, module)

Parameters:

module_name (str)
module (Module)

Return type:

bool

always_true_predicate(module_name, module)

Parameters:

module_name (str)
module (Module)

Return type:

bool

patch_module(module_name_, module)

Parameters:

module_name_ (str)
module (Module)