README.md
1# NumPy <> PyTorch Compat Layer
2
3This folder contains an implementation of (most of) the NumPy public API using PyTorch tensors.
4Note that this folder does not depend on NumPy in any way. This is a standalone implementation.
5
6This implementation is used by Dynamo to through NumPy code and lower it into PyTorch code.
7
8To see design decisions that went into this implementation, please see the [rfc](https://github.com/pytorch/rfcs/pull/54).
9
10## Structure of the code
11
12This folder exports a drop-in replacement for the NumPy namespace and its modules `linalg`, `fft` and `random` via its `__init__.py`.
13
14The implementation is split into files that work with PyTorch objects (PyTorch `Tensor`s, dtypes, etc) and files that
15use these PyTorch-only files and convert them into functions/objects that can process all the types that the NumPy functions
16accept. In particular, they accept `torch._numpy.dtype`s or `torch._numpy.ndarray`s.
17
18The PyTorch-only files are the `*_impl.py` files, while the wrapper files are those that do not have an `*_impl.py`. This creates a
19hierarchy, wherein, for example, `_dtypes.py` will import `_dtypes_impl.py`, but not the other way around. In particular, `*_impl.py`
20will only depend on other `*_impl.py` files.
21
22As discussed in the [rfc](https://github.com/pytorch/rfcs/pull/54), we use types as tags in our PyTorch implementations. We then use
23a decorator called `normalizer` that will inspect these types and preprocess the inputs before sending them to the function. This
24preprocessing is the one in charge of mapping array-like objects into `Tensor`s, dtype-like objects into PyTorch dtypes, implement
25the `out=` behaviour and so on.
26
27In the files `_funcs.py` and `_ufuncs.py` we use register the `normalizer` decorator to all the `*_impl.py` functions.
28
29In the file `_ndarray.py` we define the `ndarray` class, which is just a thin wrapper around a PyTorch tensor. We use the free functions
30and a bit of metaprogramming to implement many of the methods.
31
32## Adding a new function
33
34You just need to add a function in the relevant `*_impl.py` file. You will need to tag the inputs with the relevant Types. After that, you
35can assume that the inputs are all PyTorch objects. Your function should return PyTorch tensors. The `normalizer` will make sure that you
36always get PyTorch objects. If in doubt, you can see the implementation of the normalization attached to each type annotation in the file
37`_normalizations.py`.
38
39## Debugging
40
41It may be useful to figure out whether a given bug is caused by dynamo or the compatibility layer. You may use the compat layer in eager mode
42simply by changing `import numpy as np` by `import torch._numpy as np` in your program, without having to call `torch.compile` at all.
43Note that `torch._numpy` will be quite slow when used in eager mode, and it is in no way a replacement or an alternative to the regular PyTorch API.
44This should only be used as a debugging tool.
45