mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Type stubs for jaxlib.xla_extension no longer use -stubs suffix

Browse Source 
PEP-561 does not specify whether subpackages of a non-stub-only-package
could use the -stubs suffix. setuptools seems to allow that, yet mypy fails
to resolve the subpackage with a -stubs suffix.

This commit makes jaxlib.xla_extension a ~normal package with a toplevel
__init__.pyi.
This commit is contained in:
Sergei Lebedev 2021-08-02 14:31:11 +01:00
parent 42720f0711
commit 2a994bdb02
2 changed files with 3 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
build/build_wheel.py
View File

@ -100,13 +100,8 @@ def patch_copy_xla_extension_stubs(dst_dir):
# type stubs.
with open(os.path.join(dst_dir, "py.typed"), "w"):
pass
# The -stubs suffix is required by PEP-561.
xla_extension_dir = os.path.join(dst_dir, "xla_extension-stubs")
xla_extension_dir = os.path.join(dst_dir, "xla_extension")
os.makedirs(xla_extension_dir)
# Create a dummy __init__.py to convince setuptools that
# xla_extension-stubs is a package.
with open(os.path.join(xla_extension_dir, "__init__.py"), "w"):
pass
for stub_name in _XLA_EXTENSION_STUBS:
with open(r.Rlocation(
"org_tensorflow/tensorflow/compiler/xla/python/xla_extension/" + stub_name)) as f:

4
jaxlib/setup.py
View File

@ -30,14 +30,14 @@ setup(
description='XLA library for JAX',
author='JAX team',
author_email='jax-dev@google.com',
packages=['jaxlib', 'jaxlib.xla_extension-stubs'],
packages=['jaxlib', 'jaxlib.xla_extension'],
python_requires='>=3.7',
install_requires=['scipy', 'numpy>=1.18', 'absl-py', 'flatbuffers >= 1.12, < 3.0'],
url='https://github.com/google/jax',
license='Apache-2.0',
package_data={
'jaxlib': ['*.so', '*.pyd*', 'py.typed', 'cuda/nvvm/libdevice/libdevice*'],
'jaxlib.xla_extension-stubs': ['*.pyi'],
'jaxlib.xla_extension': ['*.pyi'],
},
zip_safe=False,
)