Merge branch 'keras-team:master' into openvino-nn-functions

Author: danielenricocahall

.github/workflows/actions.yml

@@ -148,4 +148,4 @@ jobs:
           pip uninstall -y keras keras-nightly
           pip install -e "." --progress-bar off --upgrade
       - name: Run pre-commit
-        run: pre-commit run --all-files --hook-stage manual
+        run: pre-commit run --all-files --hook-stage manual

.github/workflows/tpu_tests.yml

@@ -0,0 +1,56 @@
+name: Keras Tests
+
+# TODO: Consider enabling all tests (pytest, applications, etc.) with NNX in the future
+# Currently only basic flow tests run with NNX enabled
+
+on:
+  push:
+    branches: [ master ]
+  pull_request_review:
+    types: [submitted]
+  release:
+    types: [created]
+
+permissions:
+  contents: read
+
+jobs:
+
+  test-in-container:
+    name: Run tests on TPU
+    runs-on: linux-x86-ct6e-44-1tpu
+    # Only run on approved PRs, pushes to master, or releases
+    if: |
+      github.event_name == 'push' ||
+      github.event_name == 'release' ||
+      (github.event_name == 'pull_request_review' && github.event.review.state == 'approved')
+
+    strategy:
+      fail-fast: false
+      matrix:
+        backend: [jax]
+
+    container:
+      image: python:3.10-slim
+      options: --privileged --network host
+
+    steps:
+      - name: Checkout Repository
+        uses: actions/checkout@v4
+      
+      - name: Install Dependencies
+        run: |
+          pip install --no-cache-dir -r requirements-${{ matrix.backend }}-tpu.txt \
+      
+      - name: Set Keras Backend
+        run: echo "KERAS_BACKEND=jax" >> $GITHUB_ENV
+
+      - name: Run Verification and Tests
+        run: |
+          echo "Successfully running inside the public python container!"
+          echo "Verifying JAX installation..."
+          python3 -c "import jax; print(f'JAX backend: {jax.default_backend()}'); print(f'JAX devices : {jax.devices()}')"
+
+          pytest keras --ignore keras/src/applications \
+                      --cov=keras \
+                      --cov-config=pyproject.toml

conftest.py

@@ -31,6 +31,21 @@ def pytest_collection_modifyitems(config, items):
                 line.strip() for line in openvino_skipped_tests if line.strip()
             ]
 
+    tpu_skipped_tests = []
+    if backend() == "jax":
+        try:
+            with open(
+                "keras/src/backend/jax/excluded_tpu_tests.txt", "r"
+            ) as file:
+                tpu_skipped_tests = file.readlines()
+                # it is necessary to check if stripped line is not empty
+                # and exclude such lines
+                tpu_skipped_tests = [
+                    line.strip() for line in tpu_skipped_tests if line.strip()
+                ]
+        except FileNotFoundError:
+            pass  # File doesn't exist, no tests to skip
+
     requires_trainable_backend = pytest.mark.skipif(
         backend() in ["numpy", "openvino"],
         reason="Trainer not implemented for NumPy and OpenVINO backend.",
@@ -49,6 +64,14 @@ def pytest_collection_modifyitems(config, items):
                         "Not supported operation by openvino backend",
                     )
                 )
+        # also, skip concrete tests for TPU when using JAX backend
+        for skipped_test in tpu_skipped_tests:
+            if skipped_test in item.nodeid:
+                item.add_marker(
+                    pytest.mark.skip(
+                        reason="Known TPU test failure",
+                    )
+                )
 
 
 def skip_if_backend(given_backend, reason):

keras/api/_tf_keras/keras/ops/__init__.py

@@ -179,6 +179,7 @@
 from keras.src.ops.numpy import dot as dot
 from keras.src.ops.numpy import einsum as einsum
 from keras.src.ops.numpy import empty as empty
+from keras.src.ops.numpy import empty_like as empty_like
 from keras.src.ops.numpy import equal as equal
 from keras.src.ops.numpy import exp as exp
 from keras.src.ops.numpy import exp2 as exp2

keras/api/_tf_keras/keras/ops/numpy/__init__.py

@@ -65,6 +65,7 @@
 from keras.src.ops.numpy import dot as dot
 from keras.src.ops.numpy import einsum as einsum
 from keras.src.ops.numpy import empty as empty
+from keras.src.ops.numpy import empty_like as empty_like
 from keras.src.ops.numpy import equal as equal
 from keras.src.ops.numpy import exp as exp
 from keras.src.ops.numpy import exp2 as exp2

keras/api/ops/__init__.py

@@ -179,6 +179,7 @@
 from keras.src.ops.numpy import dot as dot
 from keras.src.ops.numpy import einsum as einsum
 from keras.src.ops.numpy import empty as empty
+from keras.src.ops.numpy import empty_like as empty_like
 from keras.src.ops.numpy import equal as equal
 from keras.src.ops.numpy import exp as exp
 from keras.src.ops.numpy import exp2 as exp2

keras/api/ops/numpy/__init__.py

@@ -65,6 +65,7 @@
 from keras.src.ops.numpy import dot as dot
 from keras.src.ops.numpy import einsum as einsum
 from keras.src.ops.numpy import empty as empty
+from keras.src.ops.numpy import empty_like as empty_like
 from keras.src.ops.numpy import equal as equal
 from keras.src.ops.numpy import exp as exp
 from keras.src.ops.numpy import exp2 as exp2

keras/src/backend/common/name_scope.py

@@ -58,7 +58,8 @@ def __exit__(self, *args, **kwargs):
             name_scope_stack = global_state.get_global_attribute(
                 "name_scope_stack"
             )
-            name_scope_stack.pop()
+            if name_scope_stack:
+                name_scope_stack.pop()
 
 
 def current_path():

keras/src/backend/common/name_scope_test.py

@@ -1,4 +1,7 @@
+import threading
+
 from keras.src import testing
+from keras.src.backend.common import global_state
 from keras.src.backend.common.name_scope import current_path
 from keras.src.backend.common.name_scope import name_scope
 
@@ -46,3 +49,85 @@ def test_override_parent(self):
                         current_path(), "absolute/path/middle/inner"
                     )
             self.assertEqual(current_path(), "outer")
+
+    def test_exit_with_none_stack(self):
+        """Test that __exit__ handles None name_scope_stack gracefully."""
+        # Create a name_scope instance
+        scope = name_scope("test")
+        # Enter the scope normally
+        scope.__enter__()
+
+        # Simulate the scenario where global state is cleared
+        # (e.g., in a different thread)
+        global_state.set_global_attribute("name_scope_stack", None)
+
+        # Exit should not raise an AttributeError
+        scope.__exit__()
+
+        # Clean up: reset the stack
+        global_state.set_global_attribute("name_scope_stack", [])
+
+    def test_exit_with_empty_stack(self):
+        """Test that __exit__ handles empty name_scope_stack gracefully."""
+        # Create a name_scope instance
+        scope = name_scope("test")
+        # Enter the scope normally
+        scope.__enter__()
+
+        # Simulate the scenario where the stack is cleared
+        name_scope_stack = global_state.get_global_attribute("name_scope_stack")
+        name_scope_stack.clear()
+
+        # Exit should not raise an IndexError
+        scope.__exit__()
+
+        # Verify stack is still empty
+        name_scope_stack = global_state.get_global_attribute(
+            "name_scope_stack", default=[]
+        )
+        self.assertEqual(len(name_scope_stack), 0)
+
+    def test_multithreaded_name_scope(self):
+        """Test name_scope in multithreaded environment."""
+        results = []
+
+        def thread_function(thread_id):
+            # Each thread should have its own name_scope_stack
+            with name_scope(f"thread_{thread_id}"):
+                path = current_path()
+                results.append(path)
+                # Verify we get the expected path
+                self.assertEqual(path, f"thread_{thread_id}")
+
+        # Create and start multiple threads
+        threads = []
+        for i in range(5):
+            thread = threading.Thread(target=thread_function, args=(i,))
+            threads.append(thread)
+            thread.start()
+
+        # Wait for all threads to complete
+        for thread in threads:
+            thread.join()
+
+        # Verify all threads executed successfully
+        self.assertEqual(len(results), 5)
+
+    def test_exit_without_pop_on_exit(self):
+        """Test that __exit__ respects _pop_on_exit flag."""
+        # Create a name_scope but don't enter it
+        scope = name_scope("test")
+        # _pop_on_exit should be False
+        self.assertFalse(scope._pop_on_exit)
+
+        # Set up a stack manually
+        global_state.set_global_attribute("name_scope_stack", [scope])
+
+        scope.__exit__()
+
+        # Verify the stack still contains the scope
+        name_scope_stack = global_state.get_global_attribute("name_scope_stack")
+        self.assertEqual(len(name_scope_stack), 1)
+
+        # Clean up
+        global_state.set_global_attribute("name_scope_stack", [])

keras/src/backend/common/variables.py

@@ -276,7 +276,7 @@ def value(self):
         return self._maybe_autocast(self._value)
 
     def assign(self, value):
-        value = self._convert_to_tensor(value, dtype=self.dtype)
+        value = self._convert_to_tensor(value, dtype=self._dtype)
         if not shape_equal(value.shape, self.shape):
             raise ValueError(
                 "The shape of the target variable and "