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

[kobihikri]

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

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.