Repository: keras-team/keras-nlp

Issue:
Fix breakages with keras-core 0.4.1
keras-core 0.4.1 released a few hours ago, and looks like it broke our CI.

https://github.com/keras-team/keras-nlp/actions/runs/5803058494/job/15730471609

We should investigate and fix.


Discussion:

Resolving patch:
diff --git a/keras_nlp/layers/modeling/masked_lm_head.py b/keras_nlp/layers/modeling/masked_lm_head.py
--- a/keras_nlp/layers/modeling/masked_lm_head.py
+++ b/keras_nlp/layers/modeling/masked_lm_head.py
@@ -141,7 +141,7 @@ def __init__(
                 )
             self.vocabulary_size = shape[0]
 
-    def build(self, inputs_shape):
+    def build(self, inputs_shape, mask_positions_shape=None):
         if self.embedding_weights is not None:
             feature_size = self.embedding_weights.shape[-1]
         else:
@@ -221,9 +221,7 @@ def get_config(self):
         )
         return config
 
-    # TODO: restore this after https://github.com/keras-team/keras-core/pull/632
-    # is in a release!
-    # def compute_output_shape(self, inputs_shape, mask_positions_shape):
-    #     output_shape = list(mask_positions_shape)
-    #     output_shape[-1] = self.vocabulary_size
-    #     return tuple(output_shape)
+    def compute_output_shape(self, inputs_shape, mask_positions_shape):
+        output_shape = list(mask_positions_shape)
+        output_shape[-1] = self.vocabulary_size
+        return tuple(output_shape)
diff --git a/keras_nlp/models/xlnet/xlnet_encoder.py b/keras_nlp/models/xlnet/xlnet_encoder.py
--- a/keras_nlp/models/xlnet/xlnet_encoder.py
+++ b/keras_nlp/models/xlnet/xlnet_encoder.py
@@ -259,30 +259,26 @@ def __init__(self, hidden_dim, kernel_initializer_range, **kwargs):
             self.kernel_initializer_range
         )
 
-    def build(self, input_shape):
+    def build(self, inputs_shape):
         self.mask_emb = self.add_weight(
             shape=(1, 1, self.hidden_dim),
             initializer=self.kernel_initializer,
             trainable=True,
             name="mask_emb",
         )
-        super().build(input_shape)
+        self.built = True
 
-    def call(
-        self,
-        padding_mask,
-        mlen=None,
-    ):
-        bsz, qlen = ops.shape(padding_mask)[0], ops.shape(padding_mask)[1]
+    def call(self, inputs, mlen=None):
+        bsz, qlen = ops.shape(inputs)[0], ops.shape(inputs)[1]
         mlen = 0 if mlen is None else mlen
 
-        padding_mask = 1 - padding_mask
-        padding_mask = ops.reshape(
-            padding_mask,
-            [ops.shape(padding_mask)[1], ops.shape(padding_mask)[0]],
+        inputs = 1 - inputs
+        inputs = ops.reshape(
+            inputs,
+            [ops.shape(inputs)[1], ops.shape(inputs)[0]],
         )
 
-        data_mask = ops.expand_dims(padding_mask, 0)
+        data_mask = ops.expand_dims(inputs, 0)
 
         if mlen > 0:
             mems_mask = ops.zeros([ops.shape(data_mask)[0], mlen, bsz])
@@ -342,14 +338,7 @@ class XLNetSegmentMatrixLayer(keras.layers.Layer):
     This layer creates Segment Matrix for XLNet Encoder.
     """
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def call(
-        self,
-        segment_ids,
-        mlen=None,
-    ):
+    def call(self, segment_ids, mlen=None):
         bsz = ops.shape(segment_ids)[0]
         mlen = 0 if mlen is None else mlen
 


Test patch: