Address comments, minor fixes

keras_nlp/models/backbone.py

@@ -29,15 +29,14 @@ def __init__(self, *args, dtype=None, **kwargs):
         )
         self._initialized = True
         if dtype is not None:
-            if isinstance(dtype, keras.mixed_precision.DTypePolicy):
-                policy = dtype
-            else:
-                policy = keras.mixed_precision.DTypePolicy(dtype)
             # Keras 2 and Keras 3 handle setting policy differently.
             if config.keras_3():
-                self.dtype_policy = policy
+                if isinstance(dtype, keras.DTypePolicy):
+                    self.dtype_policy = dtype
+                else:
+                    self.dtype_policy = keras.DTypePolicy(dtype)
             else:
-                self._dtype_policty = policy
+                self._set_dtype_policy(dtype)
 
     def __dir__(self):
         if config.keras_3():
@@ -77,7 +76,7 @@ def token_embedding(self):
 
         This layer embeds integer token ids to the hidden dim of the model.
         """
-        return self._token_embedding
+        return getattr(self, "_token_embedding", None)
 
     @token_embedding.setter
     def token_embedding(self, value):

keras_nlp/models/task.py

@@ -36,6 +36,12 @@ def __init__(self, *args, **kwargs):
             id(layer) for layer in self._flatten_layers()
         )
         self._initialized = True
+        if self.backbone is not None:
+            # Keras 2 and Keras 3 handle setting policy differently.
+            if config.keras_3():
+                self.dtype_policy = self._backbone.dtype_policy
+            else:
+                self._set_dtype_policy(self._backbone.dtype_policy)
 
     def __dir__(self):
         if config.keras_3():
@@ -128,21 +134,16 @@ def __setattr__(self, name, value):
     @property
     def backbone(self):
         """A `keras.Model` instance providing the backbone sub-model."""
-        return self._backbone
+        return getattr(self, "_backbone", None)
 
     @backbone.setter
     def backbone(self, value):
         self._backbone = value
-        # Keras 2 and Keras 3 handle setting policy differently.
-        if config.keras_3():
-            self.dtype_policy = value.dtype_policy
-        else:
-            self._dtype_policty = value.dtype_policy
 
     @property
     def preprocessor(self):
         """A `keras.layers.Layer` instance used to preprocess inputs."""
-        return self._preprocessor
+        return getattr(self, "_preprocessor", None)
 
     @preprocessor.setter
     def preprocessor(self, value):