Circular dependency between autograd.py and init.py

[navalnica]

Description

It's impossible to use needle.autograd.Tensor as type hints for functions under needle.init.

For example executing import needle as ndl in jupyter notebook fails with following error: AttributeError: module 'needle' has no attribute 'Tensor' if init.py contains following function:

import needle as ndl

def rand(*shape, low=0.0, high=1.0, device=None, dtype="float32", requires_grad=False) -> ndl.Tensor:
    """ Generate random numbers uniform between low and high """
    device = ndl.cpu() if device is None else device
    array = device.rand(*shape) * (high - low) + low
    return ndl.Tensor(array, device=device, dtype=dtype, requires_grad=requires_grad)

However, import needle as ndl works fine if there is no type hint for return value in init.py function:

import needle as ndl

def rand(*shape, low=0.0, high=1.0, device=None, dtype="float32", requires_grad=False):
    """ Generate random numbers uniform between low and high """
    device = ndl.cpu() if device is None else device
    array = device.rand(*shape) * (high - low) + low
    return ndl.Tensor(array, device=device, dtype=dtype, requires_grad=requires_grad)

The reason is that:

1. __init__.py imports autograd.Tensor before init module:

   from .autograd import Tensor, cpu, all_devices
   from . import ops
   from .ops import *
   from . import init
   from . import data
   from . import nn
   from . import optim

2. import of autograd.Tensor leads to a subsequent import of init module because autograd.py contains following line:

   from needle import init

   by the way, relative import would be much clear in this case: from . import init. but that's not the issue
3. when init.py is imported by autograd.py, function signatures inside init.py are analyzed.
   type hint def rand(...) -> ndl.Tensor leads to an error AttributeError: module 'needle' has no attribute 'Tensor'
   because the process of importing autograd in not finished yet.
   
   on import, only function signatures are analyzed. it allows to use ndl.Tensor inside function body (as is right now). the reason, I guess, is when functions are executed, all imports (including needle.autograd.Tensor) are alredy resolved.

To wrap up, it is impossible to use ndl.Tensor as type hints in init.py because of the current order of imports:
__init__.py -> autograd.py -> init.py

Solution

Solution is to:

1. change the order of imports in __init.py__ to import init earlier than autograd:

   from . import init
   from .autograd import Tensor, cpu, all_devices
   from . import ops
   from .ops import *
   from . import data
   from . import nn
   from . import optim

2. and to replace "global" import in init.py to a relative one:
   replace

   import needle as ndl

   with

   from .autograd import Tensor, cpu

[OccupyMars2025]

OccupyMars2025/Deep-Learning-Systems-Algorithms-and-Implementation#4
(2023/2/12 22:30)
Hello, I just find this issue. I don't have a good solution.