refactor rotary embeddings

DavidLandup0 · DavidLandup0 · commit a663a5c002dc · 2025-08-05T21:13:45.000+09:00
diff --git a/keras_hub/src/models/smollm3/smollm3_layers.py b/keras_hub/src/models/smollm3/smollm3_layers.py
@@ -606,7 +606,7 @@ def call(
             inv_freq_expanded, (batch_size, ops.shape(self.inv_freq)[0], 1)
         )
 
-        position_ids_expanded = ops.expand_dims(positions, axis=1)
+        position_ids_expanded = ops.expand_dims(positions, axis=1).T
 
         freqs = ops.matmul(
             ops.cast(inv_freq_expanded, "float32"),

Original file line number	Diff line number	Diff line change
`@@ -606,7 +606,7 @@ def call(`
`606`	`606`	`inv_freq_expanded, (batch_size, ops.shape(self.inv_freq)[0], 1)`
`607`	`607`	`)`
`608`	`608`
`609`		`- position_ids_expanded = ops.expand_dims(positions, axis=1)`
	`609`	`+ position_ids_expanded = ops.expand_dims(positions, axis=1).T`
`610`	`610`
`611`	`611`	`freqs = ops.matmul(`
`612`	`612`	`ops.cast(inv_freq_expanded, "float32"),`