Debug a minimal transformer computation

Question

## Scenario

A small NumPy model produces incorrect attention outputs and fails to learn. Implement corrected versions of four operations used by a decoder-style model. The model dimension is small and attention uses one head, so focus on tensor semantics rather than framework APIs.

## Interface

```python
def pos_embed(x, pos): ...
def attention(q, k, v, mask=None): ...
def linear(x, weight, bias): ...
def train_step(params, x, target, lr=0.1): ...
```

## Requirements

### `pos_embed`

Return token representations with positional embeddings added elementwise. Shapes must match or broadcast unambiguously.

### `attention`

- Compute scaled dot-product attention using a scale of `1 / sqrt(d_k)`.
- Apply a boolean mask before softmax: `True` entries are allowed and `False` entries must receive effectively zero probability.
- Use a numerically stable row-wise softmax.
- Return the weighted sum of `v`.

### `linear`

Weights have shape `(out_features, in_features)`. Return `x @ weight.T + bias`.

### `train_step`

- `params` contains `w` and `b` for a linear model.
- Compute predictions, mean squared error, gradients for both parameters, and one gradient-descent update.
- Return `(new_params, loss)` without mutating the input parameter dictionary or arrays.

Assume inputs are floating-point NumPy arrays with compatible dimensions.

## Debugging discussion

For each function, explain the observable symptom of a common defect: subtracting position, masking with zero before softmax, using the wrong weight orientation, or computing loss without updating parameters. Include at least one small invariant or test that would isolate the defect.

## Follow-ups

1. Extend attention to batches and multiple heads while documenting every shape.
2. Add a causal mask and verify that future-token values cannot affect earlier outputs.
3. Introduce a key/value cache for autoregressive decoding and explain cache growth.
4. Compare finite-difference gradient checks with automatic differentiation.