Fix shape passed to inner EinsumDense in XLNET attention. (#2751)

hertschuh · web-flow · commit 19551d66c226 · 2026-05-29T13:49:46.000-07:00
A reccent fix in Keras keras-team/keras#23000 uncovered an issue whereby a call to EinsumDense.compute_output_shape was passing an int instead of an input shape.
diff --git a/keras_hub/src/models/xlnet/relative_attention.py b/keras_hub/src/models/xlnet/relative_attention.py
@@ -192,7 +192,7 @@ def build(self, content_stream_shape):
             **self._get_common_kwargs_for_sublayer(),
         )
         self._output_dense.build(
-            self._value_dense.compute_output_shape(self._value_dim)
+            self._value_dense.compute_output_shape(self._value_shape)
         )
 
         einsum_equation, _, output_rank = _build_proj_equation(

Original file line number	Diff line number	Diff line change
`@@ -192,7 +192,7 @@ def build(self, content_stream_shape):`
`192`	`192`	`**self._get_common_kwargs_for_sublayer(),`
`193`	`193`	`)`
`194`	`194`	`self._output_dense.build(`
`195`		`- self._value_dense.compute_output_shape(self._value_dim)`
	`195`	`+ self._value_dense.compute_output_shape(self._value_shape)`
`196`	`196`	`)`
`197`	`197`
`198`	`198`	`einsum_equation, _, output_rank = _build_proj_equation(`