Remove internal_test_util folder and packages from jax wheel.

PiperOrigin-RevId: 736861450

Author: Google-ML-Automation

BUILD.bazel

@@ -31,9 +31,6 @@ transitive_py_deps(
         "//jax:experimental",
         "//jax:experimental_colocated_python",
         "//jax:experimental_sparse",
-        "//jax:internal_export_back_compat_test_util",
-        "//jax:internal_test_harnesses",
-        "//jax:internal_test_util",
         "//jax:lax_reference",
         "//jax:pallas_experimental_gpu_ops",
         "//jax:pallas_gpu_ops",

setup.py

@@ -50,7 +50,7 @@ def load_version_module(pkg_path):
     long_description_content_type='text/markdown',
     author='JAX team',
     author_email='[email protected]',
-    packages=find_packages(exclude=["examples", "jax/src/internal_test_util"]),
+    packages=find_packages(exclude=["*examples*", "*internal_test_util*"]),
     package_data={'jax': ['py.typed', "*.pyi", "**/*.pyi"]},
     python_requires='>=3.10',
     install_requires=[