Repository: huggingface/transformers
Pull request: [xLSTM] Fix bugs preventing small model training

# What does this PR do?

Fixes three bugs in the xLSTM implementation that prevent training models smaller than 7B parameters.

Fixes #43208

## Changes

1. **Line 235**: Fixed typo where tensor was called as function instead of using `.reshape()`
2. **Lines 184-185**: Fixed shape mismatch by using `dqk // nc` instead of `dqk` for proper head dimension (verified against [NX-AI/xlstm original](https://github.com/NX-AI/xlstm/blob/main/xlstm/blocks/mlstm/backends.py))
3. **Line 783**: Fixed `return_last_states` default from `bool = False` to `bool | None = None` to match docstring and enable config-based control

## Testing

Tested with 125M parameter model (hidden_size=768, num_heads=4, layers=4):

**Before fix:**
```
RuntimeError: shape '[2, 4, 2, 64, 192]' is invalid for input of size 98304
```

**After fix:**
```
Model created successfully! Parameters: 129.01M
Forward pass successful!
Output shape: torch.Size([2, 128, 65536])
```

All existing xLSTM tests pass (59 passed, 162 skipped for hardware-specific tests).

## Before submitting
- [x] This PR fixes a typo or improves the docs (if applicable)
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md)?
- [x] Was this discussed/approved via a Github issue? #43208
- [x] Did you make sure to update the documentation with your changes?
- [x] Did you write any new necessary tests?

Changed files:
filename: src/transformers/models/xlstm/modeling_xlstm.py
patch: @@ -176,12 +176,13 @@ def mlstm_chunkwise_parallel_fw_H(
         _device = matQ.device
         nc = num_chunks
         batch_size, nh, dqk, dhv = matC_states.shape
-        matC_k_states = matC_states.view(batch_size, nh, nc, dqk // nc, dhv)
-        vecN_k_states = vecN_states.view(batch_size, nh, nc, dqk // nc)
+        dhqk = dqk // nc
+        matC_k_states = matC_states.view(batch_size, nh, nc, dhqk, dhv)
+        vecN_k_states = vecN_states.view(batch_size, nh, nc, dhqk)
         scaMinter_k_states = scaMinter_states
 
-        matQ = matQ.view(batch_size, nh, nc, chunk_size, dqk)
-        matK = matK.view(batch_size, nh, nc, chunk_size, dqk)
+        matQ = matQ.view(batch_size, nh, nc, chunk_size, dhqk)
+        matK = matK.view(batch_size, nh, nc, chunk_size, dhqk)
         matV = matV.view(batch_size, nh, nc, chunk_size, dhv)
 
         ltr = torch.tril(
@@ -232,7 +233,7 @@ def mlstm_chunkwise_parallel_fw_H(
 
         # we need the denominator and the overall max state for the backward pass
         vecN_out = vecDenom_max_common.reshape(batch_size, nh, nc * chunk_size)
-        vecM_out = vecM_k_combine(batch_size, nh, nc * chunk_size)
+        vecM_out = vecM_k_combine.reshape(batch_size, nh, nc * chunk_size)
         return matH_out, vecN_out, vecM_out
 
     def mlstm_chunkwise_fw(
@@ -780,7 +781,7 @@ def forward(
             c_initial: torch.Tensor | None = None,
             n_initial: torch.Tensor | None = None,
             m_initial: torch.Tensor | None = None,
-            return_last_states: bool = False,
+            return_last_states: bool | None = None,
             mode: Literal["train", "inference"] | None = None,
         ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
             """Forward pass of the mLSTM backend.

filename: tests/models/xlstm/test_modeling_xlstm.py
patch: @@ -242,6 +242,24 @@ def recursive_check(tuple_object, dict_object):
             dict_inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=True)
             check_equivalence(model, tuple_inputs, dict_inputs, {"output_hidden_states": True})
 
+    def test_chunkwise_shape_calculation(self):
+        config = self.model_tester.get_config()
+
+        config.chunkwise_kernel = "chunkwise--native_autograd"
+
+        model = xLSTMModel(config)
+        model.to(torch_device)
+        model.train(False)
+
+        batch_size, seq_length = 2, config.chunk_size * 2
+        input_ids = ids_tensor([batch_size, seq_length], config.vocab_size)
+
+        with torch.no_grad():
+            outputs = model(input_ids)
+
+        expected_shape = (batch_size, seq_length, config.hidden_size)
+        self.assertEqual(outputs.last_hidden_state.shape, expected_shape)
+
 
 @require_torch
 @slow