PyTorch backend silently accepts shapes/activations the JAX backend rejects (three cases)

[kobihikri]

The two backends are meant to agree — convert_and_upload.py even runs a verify_parity step — but the PyTorch side accepts several inputs the JAX side rejects, and does so silently. Three concrete cases, all reproduced on b15593e4c1111ddb5f4f30dd2957df2edbaa04ca in a clean container.

1. d_model not divisible by nhead

JAX raises (tabfm/src/jax/model.py:1148-1151):

if d_model % nhead != 0:
    raise ValueError(f'd_model ({d_model}) must be divisible by nhead ({nhead})')

PyTorch just truncates (tabfm/src/pytorch/model.py:102): self.nhead, self.hd = nhead, d_model // nhead.

[torch embed_dim=7 nhead=2 (not divisible)] NO ERROR -> nhead=2 hd=3 -> nhead*hd=6 vs d_model=7

The model builds and runs while quietly dropping a dimension.

2. RoPE dimension

JAX asserts dim >= 2 (jax/model.py:250). PyTorch's RoPE.__init__ (pytorch/model.py:84-87) has no guard:

[Torch RoPE dim=1] NO ERROR -> (1,)
[Torch RoPE dim=0] NO ERROR -> (0,)
[Torch RoPE dim=2 (control)] NO ERROR -> (1,)

dim=0 yields an empty frequency buffer, so the rotation silently becomes a no-op. Note also that dim=1 and dim=2 both produce shape (1,).

3. The supported activation sets differ

JAX accepts {relu, gelu, swiglu} and raises a ValueError naming the valid options (jax/model.py:1153-1155). PyTorch's get_activation (pytorch/model.py:41-47) accepts {relu, gelu, silu} and raises a bare KeyError:

[Torch get_activation('silu')]   NO ERROR -> silu
[Torch get_activation('swiglu')] KeyError: 'swiglu'
[Torch get_activation('nosuch')] KeyError: 'nosuch'

So silu works on one backend only, and an unknown name gives a message that does not say what was expected. (swiglu is handled earlier via the self.swiglu flag, so it works in practice through MLP — it is get_activation itself that disagrees.)

None of these is urgent, and the default TabFM shapes avoid all three. They matter because the failure mode is silence: a user who mis-specifies a shape gets a subtly wrong model from PyTorch and a clear error from JAX.

Disclosure: I used an AI assistant to help find these. I ran the reproductions myself.