version 0.8 seems to miss some imports

[shaobohou]

Running with --tree --trim on this file https://github.com/deepmind/tf2jax/blob/main/tf2jax/_src/ops.py

0.7 gives the correct tree

+ tf2jax/_src/ops.py
    :: dataclasses.py
    :: functools.py
    :: typing.py
    :: absl/logging/__init__.py
    :: jax/__init__.py
    :: jax/_src/lax/control_flow/__init__.py
    :: jax/numpy/__init__.py
    :: numpy/__init__.py
    :: tensorflow/__init__.py
      tf2jax/_src/config.py
        :: contextlib.py
      tf2jax/_src/numpy_compat.py
      tf2jax/_src/xla_utils.py

but 0.8 misses tf2jax/_src/config.py

+ tf2jax/_src/ops.py
    :: dataclasses.py
    :: functools.py
    :: typing.py
    :: absl/logging/__init__.py
    :: jax/__init__.py
    :: jax/_src/lax/control_flow/__init__.py
    :: jax/numpy/__init__.py
    :: numpy/__init__.py
    :: tensorflow/__init__.py
      tf2jax/_src/__init__.py
      tf2jax/_src/numpy_compat.py
      tf2jax/_src/xla_utils.py

Any idea what is causing this?

[martindemello]

this seems to be a complicated bug. the underlying issue is that this line

importlib.util.find_spec("tf2jax._src.config")

throws an exception (not sure why, it works from ipython). that makes us think that the import is tf2jax._src.__init__ (i.e. a package import of _src rather than a module import of _src.config). before #63 we would then try falling back to just mechanically converting the relative module fullname to a filepath, which caused some issues with package imports. fixing that now filters out tf2jax._src.config altogether, but #63 fixed some more serious bugs so we can't revert it.

[shaobohou]

That sounds like a bug, no?

It fails to resolve tf2jax._src.config from ops.py, but it succeeds for tf2jax.py in the same directory, except not for the direct import but for the one imported originally in ops.py

+ tf2jax/_src/tf2jax.py
    :: collections/__init__.py
    :: inspect.py
    :: typing.py
    :: absl/logging/__init__.py
    :: jax/__init__.py
    :: jax/numpy/__init__.py
    :: numpy/__init__.py
    :: tensorflow/__init__.py
    :: tf2jax/_src/__init__.py
      tf2jax/_src/ops.py
        :: dataclasses.py
        :: functools.py
        :: jax/_src/lax/control_flow/__init__.py
          tf2jax/_src/config.py
            :: contextlib.py
          tf2jax/_src/numpy_compat.py
          tf2jax/_src/xla_utils.py
      tf2jax/_src/utils.py
    :: tree/__init__.py
    :: tensorflow/python/framework/op_def_registry.py
    :: tensorflow/python/framework/ops.py

Changing the import to import tf2jax._src.config as config also works, but feels like we shouldn't have to do this.

[martindemello]

yes, it's definitely a bug, i just can't pinpoint why it occurs. this is the traceback from importlib, and as noted above i cannot reproduce it by trying to run importlib.util.find_spec from a separate program or ipython:

Exception finding spec for tf2jax._src.config                                                                                                                                                                                               
Module 'numpy.core' has no attribute 'numerictypes'                                                                                                                                                                                         
Traceback (most recent call last):                                                                                                                                                                                                          
  File "/usr/local/google/home/mdemello/github/importlab/importlab/import_finder.py", line 102, in _resolve_import_versioned                                                                                                                
    spec = importlib.util.find_spec(name)                                                                                                                                                                                                   
  File "/usr/local/google/home/mdemello/.asdf/installs/python/3.10.4/lib/python3.10/importlib/util.py", line 94, in find_spec                                                                                                               
    parent = __import__(parent_name, fromlist=['__path__'])                                                                                                                                                                                 
  File "/usr/local/google/home/mdemello/github/tf2jax/tf2jax/__init__.py", line 21, in <module>                                                                                                                                             
    from tf2jax._src.tf2jax import AnnotatedFunction                                                                                                                                                                                        
  File "/usr/local/google/home/mdemello/github/tf2jax/tf2jax/_src/tf2jax.py", line 27, in <module>                                                                                                                                          
    import tensorflow as tf                                                                                                                                                                                                                 
  File "/usr/local/google/home/mdemello/.asdf/installs/python/3.10.4/lib/python3.10/site-packages/tensorflow/__init__.py", line 37, in <module>                                                                                             
    from tensorflow.python.tools import module_util as _module_util                                                                                                                                                                         
  File "/usr/local/google/home/mdemello/.asdf/installs/python/3.10.4/lib/python3.10/site-packages/tensorflow/python/__init__.py", line 42, in <module>                                                                                      
    from tensorflow.python import data                                                                                                                                                                                                      
  File "/usr/local/google/home/mdemello/.asdf/installs/python/3.10.4/lib/python3.10/site-packages/tensorflow/python/data/__init__.py", line 21, in <module>                                                                                 
    from tensorflow.python.data import experimental                                                                                                                                                                                         
  File "/usr/local/google/home/mdemello/.asdf/installs/python/3.10.4/lib/python3.10/site-packages/tensorflow/python/data/experimental/__init__.py", line 96, in <module>                                                                    
    from tensorflow.python.data.experimental import service                                                                                                                                                                                 
  File "/usr/local/google/home/mdemello/.asdf/installs/python/3.10.4/lib/python3.10/site-packages/tensorflow/python/data/experimental/service/__init__.py", line 419, in <module>                                                           
    from tensorflow.python.data.experimental.ops.data_service_ops import distribute                                                                                                                                                         
  File "/usr/local/google/home/mdemello/.asdf/installs/python/3.10.4/lib/python3.10/site-packages/tensorflow/python/data/experimental/ops/data_service_ops.py", line 24, in <module>                                                        
    from tensorflow.python.data.experimental.ops import compression_ops                                                                                                                                                                     
  File "/usr/local/google/home/mdemello/.asdf/installs/python/3.10.4/lib/python3.10/site-packages/tensorflow/python/data/experimental/ops/compression_ops.py", line 16, in <module>                                                         
    from tensorflow.python.data.util import structure                                                                                                                                                                                       
  File "/usr/local/google/home/mdemello/.asdf/installs/python/3.10.4/lib/python3.10/site-packages/tensorflow/python/data/util/structure.py", line 30, in <module>                                                                           
    from tensorflow.python.ops import tensor_array_ops                                                                                                                                                                                      
  File "/usr/local/google/home/mdemello/.asdf/installs/python/3.10.4/lib/python3.10/site-packages/tensorflow/python/ops/tensor_array_ops.py", line 34, in <module>                                                                          
    from tensorflow.python.ops import array_ops                                                                                                                                                                                             
  File "/usr/local/google/home/mdemello/.asdf/installs/python/3.10.4/lib/python3.10/site-packages/tensorflow/python/ops/array_ops.py", line 1564, in <module>                                                                               
    _NON_AUTOPACKABLE_TYPES = set(np.core.numerictypes.ScalarType)                                                                                                                                                                          
  File "/usr/local/google/home/mdemello/.asdf/installs/python/3.10.4/lib/python3.10/site-packages/numpy/core/__init__.py", line 161, in __getattr__                                                                                         
    raise AttributeError(f"Module {__name__!r} has no attribute {name!r}")                                                                                                                                                                  
AttributeError: Module 'numpy.core' has no attribute 'numerictypes'         

[shaobohou]

Thanks for looking into it!

For the moment, I have downgraded importlab (and pytype) in my setup, so it is not urgent.