From 33bf553148dfa9f806b89de0b9eec11497a19bc8 Mon Sep 17 00:00:00 2001
From: Francois Chollet <francois.chollet@gmail.com>
Date: Mon, 22 Apr 2024 15:43:27 -0700
Subject: [PATCH] Move _tf_keras directory to the root of the pip package.

---
 keras/src/version.py |  2 +-
 pip_build.py         | 58 ++++++++++++++++++++++++++++++++++++++------
 2 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/keras/src/version.py b/keras/src/version.py
index c0168d651da..36b1de170ae 100644
--- a/keras/src/version.py
+++ b/keras/src/version.py
@@ -1,7 +1,7 @@
 from keras.src.api_export import keras_export
 
 # Unique source of truth for the version number.
-__version__ = "3.3.0"
+__version__ = "3.3.1"
 
 
 @keras_export("keras.version")
diff --git a/pip_build.py b/pip_build.py
index 887c7119e26..ff3ddf78a6c 100644
--- a/pip_build.py
+++ b/pip_build.py
@@ -60,6 +60,57 @@ def export_version_string(version, is_nightly=False, rc_index=None):
         f.write(init_contents)
 
 
+def ignore_files(_, filenames):
+    return [f for f in filenames if f.endswith("_test.py")]
+
+
+def copy_source_to_build_directory(root_path):
+    # Copy sources (`keras/` directory and setup files) to build
+    # directory
+    os.chdir(root_path)
+    os.mkdir(build_directory)
+    shutil.copytree(
+        package, os.path.join(build_directory, package), ignore=ignore_files
+    )
+    for fname in to_copy:
+        shutil.copy(fname, os.path.join(f"{build_directory}", fname))
+    os.chdir(build_directory)
+
+
+def build(root_path, is_nightly=False, rc_index=None):
+    if os.path.exists(build_directory):
+        raise ValueError(f"Directory already exists: {build_directory}")
+
+    try:
+        copy_source_to_build_directory(root_path)
+        create_legacy_directory()
+        from keras.src.version import __version__  # noqa: E402
+
+        export_version_string(__version__, is_nightly, rc_index)
+        return build_and_save_output(root_path, __version__)
+    finally:
+        # Clean up: remove the build directory (no longer needed)
+        shutil.rmtree(build_directory)
+
+
+def create_legacy_directory():
+    shutil.move(os.path.join("keras", "api", "_tf_keras"), "keras")
+    with open(os.path.join("keras", "api", "__init__.py")) as f:
+        contents = f.read()
+        contents = contents.replace("from keras.api import _tf_keras", "")
+    with open(os.path.join("keras", "api", "__init__.py"), "w") as f:
+        f.write(contents)
+
+    with open(os.path.join("keras", "_tf_keras", "__init__.py")) as f:
+        contents = f.read()
+        contents = contents.replace(
+            "from keras.api._tf_keras import keras",
+            "from keras._tf_keras import keras",
+        )
+    with open(os.path.join("keras", "_tf_keras", "__init__.py"), "w") as f:
+        f.write(contents)
+
+
 def build_and_save_output(root_path, __version__):
     # Build the package
     os.system("python3 -m build")
@@ -85,13 +136,6 @@ def build_and_save_output(root_path, __version__):
     return whl_path
 
 
-def build(root_path, is_nightly=False, rc_index=None):
-    from keras.src.version import __version__  # noqa: E402
-
-    export_version_string(__version__, is_nightly, rc_index)
-    return build_and_save_output(root_path, __version__)
-
-
 def install_whl(whl_fpath):
     print(f"Installing wheel file: {whl_fpath}")
     os.system(f"pip3 install {whl_fpath} --force-reinstall --no-dependencies")