ERROR : Unknown op 'IsFinite'

[bhoov]

When exporting to TFJS following the guide, I get the following:

ERROR : Unknown op 'IsFinite'. File an issue at https://github.com/tensorflow/tfjs/issues so we can add it, or register a custom execution with tf.registerOp()

My trained model uses JAX, and the logsumexp function requires finite bound checking which does not seem to be built into tensorflowjs.

System information

• TFJS version 4.1.0
• Willing to contribute (YES), but I have no experience with the internals of TFJS.

Describe the feature and the current behavior/state.
This is the first issue I have encountered in trying to port my code to TFJS that I have not been able to solve. The saved model.json has the IsFinite op whenever jvp is used, e.g.,:

{"name": "PartitionedCall/jax2tf_unknown_/jvp/jit_isfinite_/IsFinite_2", "op": "IsFinite", ...}
...

Will this change the current api? How?
No

Who will benefit with this feature?
Anyone who wants to use JAX out of the box.

[bhoov]

For future reference, it turns out I was able to fix this issue by creating a custom op. It solves the bounds checking in jax.nn.logsumexp.

const IsFinite = (node) => {
    if (node.inputs) {
        return tf.logicalNot(tf.logicalOr(tf.isNaN(node.inputs[0]), tf.isInf(node.inputs[0])))
    }
    return true
}

tf.registerOp('IsFinite', IsFinite);