Implement Batched Scaled Dot-Product Attention

Question

Implement attention for batched multi-head NumPy arrays on CPU. Inputs are finite arrays with the same floating-point dtype; dimensions are positive.

```python
def attention(q, k, v, mask=None, *, return_weights=False): ...
```

Shapes are `q: (B, H, Tq, Dk)`, `k: (B, H, Tk, Dk)`, and `v: (B, H, Tk, Dv)`. Return an array of shape `(B, H, Tq, Dv)`. When `return_weights=True`, return `(output, weights)`, with weights of shape `(B, H, Tq, Tk)`.

## Requirements

- Compute scores as `q @ k.transpose(-2, -1) / sqrt(Dk)`.
- Apply softmax along the key-position dimension.
- Support an optional boolean mask broadcastable to the score tensor, where false entries cannot receive probability.
- Avoid numerical overflow by using a stable softmax or the framework's stable implementation.
- For narrow floating-point inputs, use a wider accumulator when necessary and cast returned arrays back to the input dtype. If scaled scores cannot be represented as finite float64 values, raise `ValueError` rather than returning NaNs.
- A fully masked query row has zero weights and zero output, not NaNs.
- Preserve the input floating-point dtype and do not mutate the inputs. This implementation runs on CPU; GPU/device portability is a discussion extension.
- Reject incompatible shapes, nonfloating or mismatched dtypes, nonfinite inputs, and nonboolean or nonbroadcastable masks with `ValueError`.

## Staged follow-ups

1. Implement and validate the unmasked operation.
2. Add causal and padding masks with correct broadcasting.
3. Reduce memory for long sequences without changing results beyond stated tolerance.
4. Add an incremental key/value cache and explain its shape and update rules.
5. Diagnose NaNs, incorrect scaling, transposed dimensions, and mixed-precision instability.