Describe the bug
When training a Swin Transformer model (SwinTransformerBackbone) with training=True,the model crashes at PatchMerging.call() with:
ValueError: too many values to unpack (expected 3)
Arguments received by PatchMerging.call():
• x=tf.Tensor(shape=(32, 32, 3136, 96), dtype=float32)
• H=56
• W=56
Inference (training=False) works correctly and outputs (1, 49, 768). Only training crashes.
Root cause: In swin_transformer_layers.py, DropPath.call() generates a fixed 4D random mask (batch_size, 1, 1, 1),
but the Swin Block passes a 3D token sequence(B, L, C) = (32, 3136, 96).
Under TF broadcasting, the 3D input is left-padded to(1, 32, 3136, 96) and multiplied by the 4D mask (32, 1, 1, 1), producing (32, 32, 3136, 96)
— an extra batch dimension. This corrupted 4D tensor propagates to PatchMerging,which does B, L, C = ops.shape(x) expecting 3 values but receiving 4.

To Reproduce

Minimal reproduction (runnable locally, no Colab needed):
import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import numpy as np
import keras
from keras_hub.models import SwinTransformerBackbone
from keras import layers, Model

Workaround: keras 3.15 does not export keras.ops.random in public API

if not hasattr(keras.ops, "random"):
from keras.src import ops as _s
keras.ops.random = _s.random

backbone = SwinTransformerBackbone(
embed_dim=96, depths=(2,2,6,2),
num_heads=(3,6,12,24), window_size=7
)
inputs = keras.Input(shape=(224, 224, 3))
feat = backbone(inputs)
pooled = layers.GlobalAveragePooling1D()(feat)
outputs = layers.Dense(10, activation="softmax")(pooled)
model = Model(inputs, outputs)
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")

x = np.random.randn(4, 224, 224, 3).astype(np.float32)
y = np.random.randint(0, 10, size=(4,))
model.fit(x, y, epochs=1, batch_size=4) # ← crashes here

Environment:
keras-hub 0.31.1, keras 3.15.1, tensorflow 2.20.0, Python 3.11, CPU

Expected behavior

Training should complete without error.
The DropPath mask should match the rank of the input tensor: 3D input (B, L, C) should receive a 3D mask (B, 1, 1),
not a fixed 4D mask (B, 1, 1, 1).

Suggested fix in DropPath.call():

mask_shape = (batch_size,) + (1,) * (len(x.shape) - 1)
random_tensor = keep_prob + ops.random.uniform(mask_shape)

Additional context

• The bug exists since Swin Transformer was added in v0.29.1 (PR Add Swin Transformer Backbone #2298) and remains unfixed in v0.31.1.
• No existing GitHub issue reports this problem.
• DropPath drop_prob=0.1 in the Swin-T preset, so all users training with pretrained weights will hit this.
• Related secondary issue: keras 3.14.x/3.15.x public API keras/ops/init.py
  does not export the random submodule (0 occurrences), causing
  AttributeError: module 'keras.ops' has no attribute 'random' before reaching
  the rank mismatch. This may warrant a separate fix in keras core.
• The bug is backend-specific in manifestation
  but the root cause (fixed 4D mask) is backend-agnostic; JAX and PyTorch backends have equivalent broadcasting
  rules and likely exhibit the same failure.

**Would you like to help us fix it?
Yes. I can submit a PR with the one-line fix in DropPath.call() to dynamically
generate the mask shape based on input rank, plus a unit test that runs a
training step on a small Swin-T model and verifies no shape error.