ERROR:tf2onnx.tfonnx:Unsupported ops: Counter({'_SwitchN': 15, 'Merge': 15})

[llandsmeer]

New Operator

Describe the operator

Needed to run a JAX converted ODE solver as an ONNX model

Both are control flows ops, often used together: _SwitchN ,Merge

patrick-kidger/diffrax#202

Do you know this operator be constructed using existing ONNX operators?

Maybe Identity for Switch and If and a Scan or maybe a ReduceSum for Merge?

Is this operator used by any model currently? Which

Diffrax ODE solver when translated to TF

Are you willing to contribute it? (Y/N)

No, can't work on this until February. After that, maybe

Notes

[fatcat-z]

An OP like "_SwitchN" seems not a standard TF op so we won't have a plan to support it.
There is no corresponding ONNX op mapped to "Merge", probably need to leverage other ONNX ops to implement the same logic. In short term, we won't start this work soon. And your contributions are welcome!

[llandsmeer]

What defines a standard TF op? I get that this is probably some kind of uncommon operation, but it's not outside the TF API :). It originates from the tf.switch_case() produced by tf2jax on jax.lax.switch call. At that level it quite easily translates into a lot of IF statements

[llandsmeer]

Until this is fixed; for anyone running into this issue:

It seems like you can replace

tf.switch_case(index, branches_tf)

with

tf.case([
    (tf.equal(i, index), branch)
    for i, branch in enumerate(branches_tf)
], exclusive=True)

which might be less efficient but does compile

[fatcat-z]

What defines a standard TF op? I get that this is probably some kind of uncommon operation, but it's not outside the TF API :). It originates from the tf.switch_case() produced by tf2jax on jax.lax.switch call. At that level it quite easily translates into a lot of IF statements

"SwitchN" begins with an "" which is not like others, for example "Merge". I suspect that tf.switch_case() caused this op, probably by some other uncommon actions.