JAX and PyTorch constructors disagree on identical config: same value builds on one backend, raises on the other

[kobihikri]

Following on from #92 (PyTorch constructor) — I ran the same matrix against the JAX constructor, and the two backends disagree on identical inputs, in both directions.

Both at 0, on b15593e4c1111ddb5f4f30dd2957df2edbaa04ca:

parameter	PyTorch	JAX
embed_dim	builds OK	ZeroDivisionError
max_classes	builds OK	ZeroDivisionError
row_num_cls	builds OK	ZeroDivisionError
ff_factor	builds OK	ZeroDivisionError
num frequencies	builds OK	ZeroDivisionError
col_num_blocks	builds OK	builds OK
feature_group_size	builds OK	builds OK

And the asymmetry reverses for negatives — feature_group_size=-1 builds a model on JAX, while PyTorch raises RuntimeError: zeros: Dimension size must be non-negative.

So the same config.json can produce a working model on one backend and an exception on the other, and which one is stricter depends on the parameter. Since convert_and_upload.py converts a JAX checkpoint to PyTorch and runs a parity check, I think the intent is that the two agree.

The -1 failures on the JAX side are also worth a look on their own — col_num_blocks=-1 surfaces as MLIRError: Invalid type, which is a long way from telling the caller that a block count cannot be negative.

None of this bites the default shapes. It matters if you consider the two backends interchangeable, because right now a configuration that is valid on one is not necessarily valid on the other.

I would guess the fix is the same shared validation suggested in #92 and #94, applied once and used by both constructors, rather than patching each backend separately — but that is your call on the architecture.

Disclosure: I used an AI assistant to help find this. I ran both matrices myself.