From 6e6fb10ca38534164b03ef9d62fcc17090143876 Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <jakevdp@google.com>
Date: Tue, 27 Sep 2022 12:55:20 -0700
Subject: [PATCH] setup: bundle *.pyi files with distribution

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index c3c2bbf7732c..15fb2491061d 100644
--- a/setup.py
+++ b/setup.py
@@ -61,7 +61,7 @@ def generate_proto(source):
     author='JAX team',
     author_email='jax-dev@google.com',
     packages=find_packages(exclude=["examples"]),
-    package_data={'jax': ['py.typed']},
+    package_data={'jax': ['py.typed', "*.pyi", "**/*.pyi"]},
     python_requires='>=3.7',
     install_requires=[
         'absl-py',