Fixed xla_extension-stubs packaging

setuptools does not consider a directory with an __init__.pyi a package.
This made it silently skip xla_extension-stubs when building a wheel.

See https://github.com/pypa/setuptools/blob/a4dbe3457d89cf67ee3aa571fdb149e6eb544e88/setuptools/_distutils/command/build_py.py#L183

build/build_wheel.py

@@ -93,11 +93,15 @@ def copy_file(src_file, dst_dir, dst_filename=None):
 def patch_copy_xla_extension_stubs(dst_dir):
   # This file is required by PEP-561. It marks jaxlib as package containing
   # type stubs.
-  with open(os.path.join(dst_dir, "py.typed"), "w") as f:
+  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")
   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: