def _codegen()

in diffq/ts_export.py [0:0]

• 29 lines of code
• 6 McCabe index (conditional complexity)

def _codegen(quantizer: DiffQuantizer):
    # Generates the code for the given quantizer
    module = quantizer.model.__class__.__module__
    klass = quantizer.model.__class__.__name__
    model = quantizer.model

    assert not quantizer.float16
    names = {}
    for mod_name, mod in model.named_modules():
        names[mod] = mod_name
    unpack_assigns = []

    index = 0
    for qparam in quantizer._qparams:
        mod_name = names[qparam.module]
        if mod_name == '':
            full_name = qparam.name
        else:
            full_name = mod_name + '.' + qparam.name
        full_name = _get_full_name_access(full_name)
        if qparam.other is None:
            unpack_assigns.append(UNPACK_ASSIGN.format(full_name=full_name, index=index))
            index += 1
        else:
            other_name = names[(qparam.other.module, qparam.other.name)]
            other_name = _get_full_name_access(other_name)
            unpack_assigns.append(
                UNPACK_ASSIGN_SAME.format(full_name=full_name, other_name=other_name))

    return TEMPLATE.format(
        module=module,
        klass=klass,
        unpack_assigns='\n'.join(unpack_assigns))