Function get_alias_info

torch/utils/_python_dispatch.py:294–337

(func)

# Given an OpOverload, returns schema information on it.
# This is cached for efficiency, since it can involve running torchgen
def get_alias_info(func) -> SchemaInfo:
    if func in parsed_schema_map:
        return parsed_schema_map[func]
    # For ATen ops: use torchgen (since torchscript parser doesn't handle alias annotations
    # properly for some ops that output tensorlists)
    if func.namespace == "aten":
        torchgen_schema_str = str(func._schema)
        assert torchgen_schema_str.startswith("aten::")
        # remove the aten:: namespace, which is added by the torchscript parser,
        # and torchgen doesn't know how to handle
        torchgen_schema_str = torchgen_schema_str[6:]
        import re
        # the torchscript parser ends up converting int[2]=1 into int[2]=[1, 1],
        # which torchgen chokes on.
        torchgen_schema_str = re.sub(r'=\[[0, ]+\]', '=0', torchgen_schema_str)
        torchgen_schema_str = re.sub(r'=\[[1, ]+\]', '=1', torchgen_schema_str)
        # for aten::rot90
        torchgen_schema_str = torchgen_schema_str.replace("=[0, 1]", "=[0,1]")
        torchgen_schema = torchgen.model.FunctionSchema.parse(torchgen_schema_str)
        arg_schemas = [AliasInfo(
            alias_set=set() if a.annotation is None else set(a.annotation.alias_set),
            is_write=a.annotation is not None and a.annotation.is_write,
            name=a.name,
        ) for a in torchgen_schema.arguments.flat_all]
        out_schemas = [AliasInfo(
            alias_set=set() if a.annotation is None else set(a.annotation.alias_set),
            is_write=a.annotation is not None and a.annotation.is_write,
            name=a.name,
        ) for a in torchgen_schema.returns]
    else:
        # For non-aten ops, torchgen is untested so we rely on torchscript schema parsing
        arg_schemas = [AliasInfo(
            alias_set=set() if a.alias_info is None else set(a.alias_info.before_set),
            is_write=a.alias_info is not None and a.alias_info.is_write,
            name=a.name,
        ) for a in func._schema.arguments]
        out_schemas = [AliasInfo(
            alias_set=set() if a.alias_info is None else set(a.alias_info.before_set),
            is_write=a.alias_info is not None and a.alias_info.is_write,
            name=a.name,
        ) for a in func._schema.returns]
    schema_info = SchemaInfo(args=arg_schemas, outs=out_schemas)
    parsed_schema_map[func] = schema_info
    return schema_info

def return_and_correct_aliasing(func, args, kwargs, out):
    """