Skip to content

[FORMATS] ONNX export: If node branches produce different tensor ranks, breaking static type inference #728

Description

@antimora

The ONNX models included in silero-vad (silero_vad.onnx, silero_vad_16k_op15.onnx) have If nodes where the then_branch and else_branch produce outputs with different tensor ranks. This breaks compatibility with frameworks that require static type inference, such as:

Environment

  • silero-vad version: 6.2.0
  • ONNX opset: 16

Problem Details

The silero_vad.onnx model contains 10 If nodes with mismatched output ranks:

Location then_branch rank else_branch rank
Nested depth 1 2 3
Nested depth 2 3 2
Nested depth 2 3 2
Nested depth 2 3 2
... ... ...

The pattern appears to be:

  • One branch applies Unsqueeze (producing rank N+1)
  • Other branch uses Identity (producing rank N)

Example from the decoder's LSTMCell state handling:
If condition:
then_branch: Unsqueeze → output rank 3
else_branch: Identity → output rank 2

Why This Matters

According to the ONNX specification, If node outputs should have consistent types across branches for predictable behavior. While dynamic runtimes like ONNX Runtime can handle this at execution time, static compilers cannot determine the output type at compile time.

This prevents using silero-vad with:

  1. Rust frameworks (Burn, tract) that generate statically-typed code
  2. Ahead-of-time compilation targets
  3. Embedded/edge deployment scenarios requiring static shapes

Suggested Solution

When exporting to ONNX, ensure both branches produce the same output rank. Options:

  1. Apply matching transformations in both branches:

Instead of:

if condition:
out = unsqueeze(x, 0) # rank 3
else:
out = x # rank 2 (Identity)

Use:

if condition:
out = unsqueeze(x, 0) # rank 3
else:
out = unsqueeze(x, 0) # rank 3 (same transformation)
2. Or use explicit reshape to target shape in both branches
3. Or provide an alternative "static" ONNX export that removes conditional paths (e.g., always assume state is provided)

Reproduction

import onnx
import silero_vad
import os

pkg_dir = os.path.dirname(silero_vad.__file__)
model = onnx.load(os.path.join(pkg_dir, 'data', 'silero_vad.onnx'))

def check_if_nodes(graph, path=''):
    for i, node in enumerate(graph.node):
        if node.op_type == 'If':
            then_branch = else_branch = None
            for attr in node.attribute:
                if attr.name == 'then_branch':
                    then_branch = attr.g
                elif attr.name == 'else_branch':
                    else_branch = attr.g

            if then_branch and else_branch:
                for j, (t, e) in enumerate(zip(then_branch.output, else_branch.output)):
                    t_rank = len(t.type.tensor_type.shape.dim)
                    e_rank = len(e.type.tensor_type.shape.dim)
                    if t_rank != e_rank:
                        print(f'{path}If[{i}] output {j}: then_rank={t_rank}, else_rank={e_rank}')

                check_if_nodes(then_branch, f'{path}If[{i}].then.')
                check_if_nodes(else_branch, f'{path}If[{i}].else.')

check_if_nodes(model.graph)

This affects anyone trying to use silero-vad with static ONNX compilers. Happy to help test any proposed fixes!

Originated from this task: tracel-ai/burn#1941

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions