Search code examples
pythonpytorchembedding

How can you find where a specified function is defined in a library?


For example, I am investigating the implementation of torch.nn.Embedding class. I guessed this would be in some file in nn directory and found the class in torch/nn/module/sparse.py.

In this class, functional.embedding() function is provoked and I found this function in torch/nn/functional.py.

In this function, torch.embedding() is provoked. Now I have totally no idea of where it is defined.

Can you tell me the way to find it, for it doesn't seem that there is embedding.py in torch directory.

Can you teach me some tips for finding specific function definition?

P.S. I havefound one in torch/jit/_shape_functions.py, but not sure at all if this is the one I am looking for.


Solution

  • Your title asks about finding a function definition "in Python".

    In code, this is perhaps what you want:

    import torch
    import source
    
    print(inspect.getsourcefile(torch.nn.Embedding))
    source, start_line = inspect.getsourcelines(torch.nn.Embedding)
    print(start_line, start_line + len(source))
    

    Output (for PyTorch 2.0.1):

    C:\Users\[redacted]\pytorch_test\Lib\site-packages\torch\nn\modules\sparse.py
    13 224
    

    However, most Python IDEs make this a lot easier. For example, if you're using PyCharm, simply holding down Ctrl and clicking the name of the class in your code will take you to the definition in the source of the library.

    You said a function functional.embedding() was 'provoked', I assume you meant to say 'invoked' or 'called'. However, looking at the source, it's not clear what function you are referring to. If you are more specific about the line of code you were looking at, it shouldn't be hard to find - but you'd probably want to share the exact version of the package you're using, a line number, and a copy of the relevant line to be sure.

    The tricky part may be resolving the name in context of the namespace where it would be called. This is why it's so much easier to use an IDE for this task, since it does all that for you continuously in the background - it knows what is available at any point in the source, and how a name would be resolved at runtime.

    For example:

    • looking at torch.nn.Embedding, you find lines 13-224 of sparse.py
    • in that class, in the method forward() (which appears to be where you're looking), Torch 2.0.1 calls embedding() on F on line 162
    • that function is to be found on lines 2109-2210 of functional.py, where it calls torch.embedding() on line 2210 in the return statement
    • that function is to be found on line 821 of _VariableFunctions.pyi - a Python interface definition file.

    This is where the trail goes cold - the last reference shows that the actual work is happening outside Python. You'd have to dig into the C++ code for the library to find the source of the function you're looking for.