Rollback of "Add CUDNN custom call for LSTM. Exposed as jax.experimen…

…tal.rnn module." PiperOrigin-RevId: 490499003
justinjfu · Nov 23, 2022 · d1fbdbc · d1fbdbc
1 parent fe56a19
commit d1fbdbc
Show file tree

Hide file tree

Showing 15 changed files with 0 additions and 1,378 deletions.
diff --git a/jax/BUILD b/jax/BUILD
@@ -248,10 +248,3 @@ pytype_library(
         ":jax",
     ],
 )
-
-pytype_library(
-    name = "rnn",
-    srcs = ["experimental/rnn.py"],
-    visibility = ["//visibility:public"],
-    deps = [":jax"],
-)
diff --git a/jax/_src/lib/__init__.py b/jax/_src/lib/__init__.py
@@ -105,9 +105,6 @@ def _parse_version(v: str) -> Tuple[int, ...]:
 # branch on the Jax github.
 xla_extension_version = getattr(xla_client, '_version', 0)
 
-if xla_extension_version > 108:
-  import jaxlib.gpu_rnn as gpu_rnn  # pytype: disable=import-error
-
 can_execute_with_token = (
     xla_extension_version >= 89 and hasattr(
         xla_client.LoadedExecutable  # type: ignore