pytype breaks on macOS

The following issue relates only to the development of Flax, not the usage. In essence, it is a pytype issue (Pytype fails on MacOSX(Mojave-10.14.9) when inferring types.), but it negatively impacts the development of Flax on macOS.

Problem you have encountered:

Upon following the instructions from contributing.md, ./tests/run_all_tests.sh breaks.

• more specifically: pytype flax/
  • even more:

pytype-single --imports_info .pytype/imports/flax.core.scope.imports --module-name flax.core.scope -V 3.8 -o .pytype/pyi/flax/core/scope.pyi --analyze-annotated --nofail --quick flax/core/scope.py

Stack trace:

ERROR:pytype.imports_map_loader Invalid imports_map entries (checking from root dir: /Users/vivelev/Desktop/flax)
ERROR:pytype.imports_map_loader   file does not exist: '/Users/vivelev/Desktop/flax/.pytype/pyi/flax/core/scope.pyi' (mapped from 'flax/core/scope')
Traceback (most recent call last):
  File "/Users/vivelev/Desktop/flax/.env/bin/pytype-single", line 8, in <module>
    sys.exit(main())
  File "/Users/vivelev/Desktop/flax/.env/lib/python3.8/site-packages/pytype/single.py", line 63, in main
    options = config.Options(sys.argv[1:], command_line=True)
  File "/Users/vivelev/Desktop/flax/.env/lib/python3.8/site-packages/pytype/config.py", line 68, in __init__
    Postprocessor(names, options, self).process()
  File "/Users/vivelev/Desktop/flax/.env/lib/python3.8/site-packages/pytype/config.py", line 406, in process
    node.processor(value)
  File "/Users/vivelev/Desktop/flax/.env/lib/python3.8/site-packages/pytype/config.py", line 541, in _store_imports_map
    self.output_options.imports_map = imports_map_loader.build_imports_map(
  File "/Users/vivelev/Desktop/flax/.env/lib/python3.8/site-packages/pytype/imports_map_loader.py", line 91, in build_imports_map
    raise ValueError(msg)
ValueError: Invalid imports_map: .pytype/imports/flax.core.scope.imports
Bad entries:
  flax/core/scope -> /Users/vivelev/Desktop/flax/.pytype/pyi/flax/core/scope.pyi

Steps to reproduce:

Information about my system:

• OS: macOS 11.0.1 (Although, judging by the pytype isssue linked above, it will probably break on older versions too.)
• Python: v3.8.6 (also tried on v3.6.12 and v3.7.9)

Steps: contributing.md

A hacky solution:

The problem seems to stem from the fact that in several places in the Flax source code the Scope class is imported as:

from flax.core import Scope

rather then:

from flax.core.scope import Scope

As a result, for reasons unknown to me, pytype on macOS thinks there is a file flax/core/Scope.py, searches for it, and fails.

Sadly, however, simply replacing all from flax.core import Scope with from flax.core.scope import Scope doesn’t work: it fixes the Flax error, but for some reason it fails to process the JAX API accurately, raising the new error No attribute 'xla' on module 'jax.interpreters'.

The only solution I found so far is as follows:

1. Replace from flax.core import Scope with from flax.core.scope import Scope in the following files:

• flax/core/nn/attention.py
• flax/core/nn/linear.py
• flax/core/nn/normalization.py

1. First run pytype flax/core/scope.py and then pytype flax/

(If you want, I can open a PR with the fixes. Although, imo they are really brittle.)

A proposal:

I see you want to run both mypy and pytype in run_all_tests.sh (#688), but why not only mypy? From my experience, mypy is better supported for OSs other than Linux. Is there a specific benefit you seek from using pytype? The inference? Not being as strict as mypy?

I would be more than glad to help you transition from pytype to mypy. 😃

[jheek]

A PR is very much welcome actually! Here is the “PR welcome” issue tracking this: https://github.com/google/flax/issues/688

[avital]

pytype and mypy differ in various fine details, but we need pytype to pass for the internal Google build system. Looks like Google-internal pytype and open source pytype also differ in some ways, though.