|
| 1 | +# NumPy <> PyTorch Compat Layer |
| 2 | + |
| 3 | +This folder contains an implementation of (most of) the NumPy public API using PyTorch tensors. |
| 4 | +Note that this folder does not depend on NumPy in any way. This is a standalone implementation. |
| 5 | + |
| 6 | +This implementation is used by Dynamo to through NumPy code and lower it into PyTorch code. |
| 7 | + |
| 8 | +To 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 | + |
| 12 | +This folder exports a drop-in replacement for the NumPy namespace and its modules `linalg`, `fft` and `random` via its `__init__.py`. |
| 13 | + |
| 14 | +The implementation is split into files that work with PyTorch objects (PyTorch `Tensor`s, dtypes, etc) and files that |
| 15 | +use these PyTorch-only files and convert them into functions/objects that can process all the types that the NumPy functions |
| 16 | +accept. In particular, they accept `torch._numpy.dtype`s or `torch._numpy.ndarray`s. |
| 17 | + |
| 18 | +The PyTorch-only files are the `*_impl.py` files, while the wrapper files are those that do not have an `*_impl.py`. This creates a |
| 19 | +hierarchy, wherein, for example, `_dtypes.py` will import `_dtypes_impl.py`, but not the other way around. In particular, `*_impl.py` |
| 20 | +will only depend on other `*_impl.py` files. |
| 21 | + |
| 22 | +As discussed in the [rfc](https://github.com/pytorch/rfcs/pull/54), we use types as tags in our PyTorch implementations. We then use |
| 23 | +a decorator called `normalizer` that will inspect these types and preprocess the inputs before sending them to the function. This |
| 24 | +preprocessing is the one in charge of mapping array-like objects into `Tensor`s, dtype-like objects into PyTorch dtypes, implement |
| 25 | +the `out=` behaviour and so on. |
| 26 | + |
| 27 | +In the files `_funcs.py` and `_ufuncs.py` we use register the `normalizer` decorator to all the `*_impl.py` functions. |
| 28 | + |
| 29 | +In the file `_ndarray.py` we define the `ndarray` class, which is just a thin wrapper around a PyTorch tensor. We use the free functions |
| 30 | +and a bit of metaprogramming to implement many of the methods. |
| 31 | + |
| 32 | +## Adding a new function |
| 33 | + |
| 34 | +You 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 |
| 35 | +can assume that the inputs are all PyTorch objects. Your function should return PyTorch tensors. The `normalizer` will make sure that you |
| 36 | +always 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 | + |
| 41 | +It 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 |
| 42 | +simply by changing `import numpy as np` by `import torch._numpy as np` in your program, without having to call `torch.compile` at all. |
| 43 | +Note 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. |
| 44 | +This should only be used as a debugging tool. |
0 commit comments