dmx.compressor.fx.tracer.DmxHFTracer

class dmx.compressor.fx.tracer.DmxHFTracer(autowrap_modules=(<module 'math' from '/home/docs/.asdf/installs/python/3.11.12/lib/python3.11/lib-dynload/math.cpython-311-x86_64-linux-gnu.so'>, ), autowrap_functions=(<function apply_rotary_pos_emb>, <function apply_rotary_pos_emb>, <function apply_rotary_pos_emb>))

Custom HFTracer where definition of leaf nodes

__init__(autowrap_modules=(<module 'math' from '/home/docs/.asdf/installs/python/3.11.12/lib/python3.11/lib-dynload/math.cpython-311-x86_64-linux-gnu.so'>, ), autowrap_functions=(<function apply_rotary_pos_emb>, <function apply_rotary_pos_emb>, <function apply_rotary_pos_emb>))

Construct a Tracer object.

Parameters:

autowrap_modules (Tuple[ModuleType]) – defaults to (math, ), Python modules whose functions should be wrapped automatically without needing to use fx.wrap(). Backward-compatibility for this parameter is guaranteed.
autowrap_functions (Tuple[Callable, ...]) – defaults to (), Python functions that should be wrapped automatically without needing to use fx.wrap(). Backward compatibility for this parameter is guaranteed.
param_shapes_constant (bool) – When this flag is set, calls to shape, size and a few other shape like attributes of a module’s parameter will be evaluated directly, rather than returning a new Proxy value for an attribute access. Backward compatibility for this parameter is guaranteed.

Note

Backwards-compatibility for this API is guaranteed.

Methods

`__init__`([autowrap_modules, autowrap_functions])	Construct a Tracer object.
`call_module`(m, forward, args, kwargs)	Method that specifies the behavior of this `Tracer` when it encounters a call to an `nn.Module` instance.
`create_arg`(a)	A method to specify the behavior of tracing when preparing values to be used as arguments to nodes in the `Graph`.
`create_args_for_root`(root_fn, is_module[, ...])	Create `placeholder` nodes corresponding to the signature of the `root` Module.
`create_node`(kind, target, args, kwargs[, ...])	Inserts a graph node given target, args, kwargs, and name.
`create_proxy`(kind, target, args, kwargs[, ...])	Create a Node from the given arguments, then return the Node wrapped in a Proxy object.
`get_fresh_qualname`(prefix)	Gets a fresh name for a prefix and returns it.
`getattr`(attr, attr_val, parameter_proxy_cache)	Method that specifies the behavior of this `Tracer` when we call getattr on a call to an `nn.Module` instance.
`is_leaf_module`(m, module_qualified_name)	A method to specify whether a given `nn.Module` is a "leaf" module.
`iter`(obj)	Called when a proxy object is being iterated over, such as
`keys`(obj)	Called when a proxy object is has the keys() method called.
`patch_for_tracing`(root)
`path_of_module`(mod)	Helper method to find the qualified name of mod in the Module hierarchy of root.
`proxy`(node)
`to_bool`(obj)	Called when a proxy object is being converted to a boolean, such as
`trace`(root[, concrete_args, dummy_inputs, ...])	Traces root and returns the corresponding FX torch.fx.Graph representation.

Attributes

`allow_insert_stateless_mods`
`check_mutable_operations`
`proxy_buffer_attributes`
`record_stack_traces`
`supported_archs`
`trace_asserts`
`traced_func_name`
`graph`
`scope`
`module_stack`
`node_name_to_scope`