refactor: use dataclass for weight axis defaults table#30
Conversation
Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
| } | ||
|
|
||
| # Maps aten OpOverload -> nn.Module type for graph-mode op to module type resolution. | ||
| _ATEN_OP_TO_MODULE_TYPE: dict[torch._ops.OpOverload, type[nn.Module]] = { |
There was a problem hiding this comment.
This mapping in particular doesn't seem specific to axis defaults but more general to torch. Can we move it to torch_utils.py and have it define the mapping explicitly instead of pulling from _WEIGHT_AXIS_SPECS (which doesn't need to own the mapping)
There was a problem hiding this comment.
makes sense!
My intention here was that when we add a new entry for defaults, we only change it in one place. But with a separate table for the aten to module mapping, it will be 2 places.. I think that's fine since we get the mapping as a util of its own, separate from the defaults.
There was a problem hiding this comment.
I see, yeah I feel like the table is generic enough that other things may also make use of it, feels better to separate out the responsibilities in this case.
|
|
||
| if isinstance(granularity, PerChannelGranularity): | ||
| expected_axis = _PER_CHANNEL_WEIGHT_AXIS_DEFAULTS[module_type] | ||
| expected_axis = _WEIGHT_AXIS_SPECS[module_type].per_channel_axis |
There was a problem hiding this comment.
Can we also test the default_axis_for function here?
Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
2 participants
Use a
dataclassinstead of simple dict for the axis defaults table.