javierdejesusda opened a new pull request, #19772:
URL: https://github.com/apache/tvm/pull/19772
### Root cause
In the ONNX `LayerNormalization` spec the bias `B` is optional; when omitted
it should behave as
zeros shaped and typed like the scale `W`. In
`LayerNormalization._impl_v17`, the synthesized zero
bias instead took its shape from `data.struct_info.shape[1]` (an unrelated
data dim) and hardcoded
`dtype="float32"`. For input `[2, 3, 4, 8]` with scale `[8]` and `axis=-1`
this builds a bias of
shape `(3,)` while gamma is `(8,)`, so `relax.op.nn.layer_norm` raises a
size-mismatch
`InternalError`. The float32 hardcode also breaks fp16/bf16 no-bias models,
since gamma, beta, and
data must share a dtype. PyTorch's `nn.LayerNorm(..., bias=False)` exports
exactly this no-bias form.
### Fix
Derive both the shape and dtype of the synthesized zero bias from the scale,
matching the ONNX
semantics for an omitted `B` and the existing torch frontend
(`relax.const(np.zeros(shape), x.struct_info.dtype)`):
```python
if bias is None:
bias = relax.const(_np.zeros(gamma_shape, dtype=scale.struct_info.dtype))
```
`gamma_shape` and the `_np`/`get_const_tuple` imports are already present.
Deriving the dtype from
the scale (rather than the issue's float32-only suggestion) is what also
fixes the fp16/bf16 case.
### Test plan
Added non-square no-bias regression cases to
`test_frontend_onnx.py::test_layer_norm` (the previous
no-bias case was square, which masked the bug): float32 `[2,3,4,8]`/scale
`[8]` and float16 with
full `check_correctness`, plus a bf16 importer-only case (ORT's CPU provider
has no bf16
LayerNormalization kernel).
Fixes #19691
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
