Remove deprecated jax.experimental.stax

The new location is jax.example_libraries.stax

CHANGELOG.md

@@ -18,6 +18,8 @@ PLEASE REMEMBER TO CHANGE THE '..main' WITH AN ACTUAL TAG in GITHUB LINK.
     for an alternative API. 
   * {func}`jax.tree_util.tree_multimap` has been removed. It has been deprecated since
     JAX release 0.3.5, and {func}`jax.tree_util.tree_map` is a direct replacement.
+  * Removed `jax.experimental.stax`; it has long been a deprecated alias of
+    {mod}`jax.example_libraries.stax`.
 
 ## jaxlib 0.3.16 (Unreleased)
 * [GitHub commits](https://github.com/google/jax/compare/jaxlib-v0.3.15...main).

jax/BUILD

@@ -144,7 +144,6 @@ pytype_library(
     name = "stax",
     srcs = [
         "example_libraries/stax.py",
-        "experimental/stax.py",
     ],
     visibility = ["//visibility:public"],
     deps = [":jax"],

jax/experimental/stax.py

@@ -1,30 +0,0 @@
-# Copyright 2021 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Stax has moved to jax.example_libraries.stax
-
-jax.experimental.stax is deprecated and will delegate to
-jax.example_libraries.stax with a warning for backwards-compatibility
-for a limited time.
-"""
-
-import warnings
-
-from jax.example_libraries.stax import *    # noqa: F401,F403
-
-_HAS_DYNAMIC_ATTRIBUTES = True
-
-warnings.warn('jax.experimental.stax is deprecated, '
-              'import jax.example_libraries.stax instead',
-              FutureWarning)