From b0a318c25f99aad9598d3e654d341f3b7fc8a191 Mon Sep 17 00:00:00 2001
From: Hongyu Chiu <20734616+james77777778@users.noreply.github.com>
Date: Sun, 21 Jul 2024 19:43:50 +0800
Subject: [PATCH] Fix CI

---
 keras/src/backend/jax/trainer.py   | 4 ++--
 keras/src/backend/numpy/trainer.py | 8 ++++----
 keras/src/backend/torch/trainer.py | 4 ++--
 keras/src/layers/layer.py          | 2 ++
 keras/src/trainers/trainer.py      | 5 +++--
 5 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/keras/src/backend/jax/trainer.py b/keras/src/backend/jax/trainer.py
index cab8333688ab..ab3d49e6a2cd 100644
--- a/keras/src/backend/jax/trainer.py
+++ b/keras/src/backend/jax/trainer.py
@@ -536,7 +536,7 @@ def evaluate(
                 steps_per_execution=self.steps_per_execution,
             )
 
-        self._symbolic_build(iterator=epoch_iterator)
+        self._symbolic_build(iterator=epoch_iterator, training=False)
 
         # Container that configures and calls callbacks.
         if not isinstance(callbacks, callbacks_module.CallbackList):
@@ -765,7 +765,7 @@ def test_on_batch(
         data = (x, y, sample_weight)
         data = _distribute_data(data)
         # Maybe build model
-        self._symbolic_build(data_batch=data)
+        self._symbolic_build(data_batch=data, training=False)
         self._record_training_state_sharding_spec()
         self.make_test_function()
 
diff --git a/keras/src/backend/numpy/trainer.py b/keras/src/backend/numpy/trainer.py
index 6d40982be43e..553602363e5c 100644
--- a/keras/src/backend/numpy/trainer.py
+++ b/keras/src/backend/numpy/trainer.py
@@ -91,7 +91,7 @@ def multi_predict_steps(data):
 
         self.predict_function = predict_step
 
-    def _symbolic_build(self, data_batch):
+    def _symbolic_build(self, data_batch, training=True):
         model_unbuilt = not all(layer.built for layer in self._flatten_layers())
         compile_metrics_unbuilt = (
             self._compile_metrics is not None
@@ -113,7 +113,7 @@ def to_symbolic_input(v):
             ) = data_adapter_utils.unpack_x_y_sample_weight(data_batch)
             # Build all model state with `backend.compute_output_spec`.
             try:
-                y_pred = backend.compute_output_spec(self, x)
+                y_pred = backend.compute_output_spec(self, x, training=training)
             except:
                 raise RuntimeError(
                     "Unable to automatically build the model. "
@@ -246,7 +246,7 @@ def evaluate(
             # Build the model on one batch of data.
             for _, data in epoch_iterator.enumerate_epoch():
                 data_batch = data[0]
-                self._symbolic_build(data_batch)
+                self._symbolic_build(data_batch, training=False)
                 break
 
         # Container that configures and calls callbacks.
@@ -304,7 +304,7 @@ def test_on_batch(
         data = (x, y, sample_weight)
 
         # Maybe build model
-        self._symbolic_build(data)
+        self._symbolic_build(data, training=False)
         self.make_test_function()
 
         logs = self.test_function([data])
diff --git a/keras/src/backend/torch/trainer.py b/keras/src/backend/torch/trainer.py
index bf0e133d67d1..017d8d786b96 100644
--- a/keras/src/backend/torch/trainer.py
+++ b/keras/src/backend/torch/trainer.py
@@ -344,7 +344,7 @@ def evaluate(
                 steps_per_execution=self.steps_per_execution,
             )
 
-        self._symbolic_build(iterator=epoch_iterator)
+        self._symbolic_build(iterator=epoch_iterator, training=False)
 
         # Container that configures and calls callbacks.
         if not isinstance(callbacks, callbacks_module.CallbackList):
@@ -483,7 +483,7 @@ def test_on_batch(
         data = (x, y, sample_weight)
 
         # Maybe build model
-        self._symbolic_build(data_batch=data)
+        self._symbolic_build(data_batch=data, training=False)
         self.make_test_function()
 
         logs = self.test_function([data])
diff --git a/keras/src/layers/layer.py b/keras/src/layers/layer.py
index dc945e227813..68ba4e1e06b9 100644
--- a/keras/src/layers/layer.py
+++ b/keras/src/layers/layer.py
@@ -1143,6 +1143,8 @@ def _get_regularization_losses(self):
                 v = backend.get_stateless_scope().get_current_value(variable)
             else:
                 v = variable
+            if v is None:
+                v = variable
             weight_regularization_losses.append(variable.regularizer(v))
         return weight_regularization_losses
 
diff --git a/keras/src/trainers/trainer.py b/keras/src/trainers/trainer.py
index 8cf4cd9445b6..4a50edb81cab 100644
--- a/keras/src/trainers/trainer.py
+++ b/keras/src/trainers/trainer.py
@@ -1000,7 +1000,7 @@ def _assert_compile_called(self, method_name=None):
                 msg += f"calling `{method_name}()`."
             raise ValueError(msg)
 
-    def _symbolic_build(self, iterator=None, data_batch=None):
+    def _symbolic_build(self, iterator=None, data_batch=None, training=True):
         model_unbuilt = not all(layer.built for layer in self._flatten_layers())
         compile_metrics_unbuilt = (
             self._compile_metrics is not None
@@ -1035,7 +1035,7 @@ def to_symbolic_input(v):
 
             # Build all model state with `backend.compute_output_spec`.
             try:
-                y_pred = backend.compute_output_spec(self, x)
+                y_pred = backend.compute_output_spec(self, x, training=training)
             except Exception as e:
                 raise RuntimeError(
                     "Unable to automatically build the model. "
@@ -1065,6 +1065,7 @@ def to_symbolic_input(v):
                     y,
                     y_pred,
                     sample_weight=sample_weight,
+                    training=training,
                 )
         if optimizer_unbuilt:
             # Build optimizer