Integer floor for negative input yields wrong result

[Gerstenberger]

Describe the bug
Currently, we need a ceildiv operation to compute the sequence lengths after convolution for sequence models.
The ceildiv operation is implemented as

ceildiv(a, b) := -floordiv(-a, b)

For integer inputs, after conversion from TensorFlow to ONNX, this yields the wrong results. Looks like it is doing a floordiv(a, b) instead.
For floating point inputs (changing the signature types), the programs below yield the correct expected result.
Please see example with outputs below in to reproduce

Urgency

Currently we have this case in one of our production models, which yields wrong recognition due to the bug.

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 18.04*): Linux Ubuntu 18.04.6 LTS
• TensorFlow Version: 2.12.0
• Python version: 3.9.13
• ONNX version: 1.14.0
• ONNXRuntime version: 1.14.1

To Reproduce
Running the following program

import tensorflow as tf

x = [3, 7, 10]

class CustomModule(tf.Module):

    @tf.function(input_signature=[tf.TensorSpec([None], tf.int32), tf.TensorSpec([], tf.int32)])
    def __call__(self, a, b):
        return -tf.math.floordiv(-a, b)

ceildiv = CustomModule()

# correctly produces [1, 3, 4]
print(ceildiv(x, 3).numpy())

ceildiv(x, 3)
tf.saved_model.save(ceildiv, "saved-model-example")

then exporting it to onnx with

python3 -m tf2onnx.convert --saved-model saved-model-example/ --output example.onnx

and running it with

import numpy as np
import onnxruntime as ort

ort_sess = ort.InferenceSession('example.onnx')

x = np.array([3, 7, 10]).astype(np.int32)
y = np.array([3]).astype(np.int32)

print(ort_sess.run(None, {'a': x, 'b': y}))

yields the wrong behaviour.

For me the output is

First Script: [1, 3, 4]
Second Script: [1, 2, 3]
Second Script floats: [1.0, 3.0, 4.0]

Screenshots
Graph for Integer:

Graph for floats:

Additional context

None

[fatcat-z]

The behavior of tf.math.floordiv() is different than ONNX op floor() which is used to implement similar function of tf.math.floordiv().

The difference is how to hand integral value:
tf.math.floordiv():

rounding toward the most negative integer.

onnx op floor():

If x is integral, +0, -0, NaN, or infinite, x itself is returned.

So the attached onnx file where there is no floor op is expected. Till now, there is no good solution for solving this problem yet. We have some thoughts but need some time to try.

[albertz]

What workaround do you suggest? Maybe sth like:

def onnx_compat_floor_div(a, b):
  # https://github.com/onnx/tensorflow-onnx/issues/2174
  abs_a, abs_b = tf.abs(a), tf.abs(b)
  return tf.where(
    a * b >= 0,
    a // b,
    -abs_a // abs_b - tf.cast(abs_a % abs_b != 0, dtype=a.dtype)
  )

?

[fatcat-z]

What workaround do you suggest? Maybe sth like:

def onnx_compat_floor_div(a, b):
  # https://github.com/onnx/tensorflow-onnx/issues/2174
  abs_a, abs_b = tf.abs(a), tf.abs(b)
  return tf.where(
    a * b >= 0,
    a // b,
    -abs_a // abs_b - tf.cast(abs_a % abs_b != 0, dtype=a.dtype)
  )

?

This is certainly a possible workaround for such case, while we are looking for a way to address it within tf2onnx to avoid such unnecessary work from users.