JAX RotaryEmbedding: assert admits dim==2, then divides by (dim-2) -> ZeroDivisionError

`RotaryEmbedding.__init__` guards `dim >= 2`, and then the next statement divides by `dim - 2`:

```python
assert dim >=2, f'dim must be at least 2. Got {dim}'
# Apply theta rescaling based on NTK-aware scaling for longer sequence lengths
theta *= theta_rescale_factor ** (dim / (dim - 2))
```

`tabfm/src/jax/model.py:250-252`

So the guard admits exactly the one value the following line cannot handle. `dim == 2` raises `ZeroDivisionError` before anything else happens.

Reproduced on `b15593e4c1111ddb5f4f30dd2957df2edbaa04ca`, in a clean container with `pip install -e ".[jax,pytorch]"`:

```
[JAX RotaryEmbedding dim=2] ZeroDivisionError: division by zero
[JAX RotaryEmbedding dim=3 (control)] NO ERROR -> RotaryEmbedding
[JAX RotaryEmbedding dim=1 (assert should fire)] AssertionError: dim must be at least 2. Got 1
```

`dim` here is the per-head dimension, `d_model // nhead`, so any configuration whose head dimension works out to 2 hits this — for example `d_model=4, nhead=2`. The default TabFM shapes do not, which is presumably why it has not been seen.

Note that the exponent is evaluated regardless of `theta_rescale_factor`, so the default `theta_rescale_factor=1` does not avoid it even though `1 ** x` is 1 for every finite `x`.

Two ways to fix it, and I did not want to pick for you:

1. Tighten the guard to `dim > 2` if a head dimension of 2 is genuinely unsupported.
2. Skip the rescale when `theta_rescale_factor == 1`. That is behaviour-preserving — the term is exactly 1 in that case — and it makes `dim == 2` work on the default path.

Happy to send a PR for whichever you prefer.

Disclosure: I used an AI assistant to help find this. I ran the reproduction myself.

Start at tabfm/src/jax/model.py:250-252 and run the reported JAX dim=2, dim=3, and dim=1 reproduction after installing the JAX dependencies. Resolve whether dim=2 should be supported or rejected, then verify that the selected behavior avoids the ZeroDivisionError while preserving the existing dim=1 assertion.

Written by the indexing model from the issue text.