[AHMETHAKANBEZIR1]

run_sgd(..., shuffle=True) silently disables shuffling when the batch size exceeds the number of batches, even if there are multiple batches. For example, 10 samples with batch_size=5 always use the original order. This also affects SSM.fit_sgd, which delegates its updates to this utility.

Compare the batch size with the number of samples instead. Keep the existing no-shuffle shortcut for a full-data batch, and preserve minibatch/epoch/loss definitions.

Validation on Windows / Python 3.12.14 / JAX 0.12.0, CPU:

• On unchanged main: 2 failed / 6 passing controls. All eight pass with this fix.
• Three epochs are compared with explicit permutations and analytic scalar SGD updates, checking parameters and losses. Cases include a final short batch, batch_size=1, a full-data batch, and shuffle=False.
• Full pytest --cov=./: 245 passed (736 warnings), 81% coverage.
• Configured docstring gate and git diff --check pass.

Prepared autonomously with OpenAI Codex, including code and tests; no independent human review is claimed.