Better ONNX export for aten_floor_divide in onnx script

[bas-aarts]

The ONNX export for aten_floor_divide contains the following code

onnxscript/onnxscript/function_libs/torch_lib/ops/core.py

Lines 4021 to 4028 in e6fad72

	# Convert truncation to flooring
	# Reference: https://github.com/pytorch/pytorch/blob/ffc645c870f0abd368606ba1e2b3b58cacb03046/torch/_refs/__init__.py#L1401C1-L1409C70
	# offset = (torch.signbit(a) != torch.signbit(b)).logical_and(torch.fmod(a, b) != 0)
	# return prims.div(a, b) - _maybe_convert_to_dtype(offset, a.dtype)
	offset = op.And(
	op.Not(op.Equal(op.Sign(self), op.Sign(other))),
	op.Cast(op.Mod(self, other), to=BOOL.dtype),
	)

The expression !(sign(a) == sign(b)) && (bool)(a%b) can be improved upon

• it has 7 operators (including 2 Signs, which have limited integral support for some EPs, at least per Gemma)
• for a common case where b is a known positive integer, this expression simply optimizes to !(sign(a) == 1) && (bool)(a%b). If a is an expression that is not negative, this is tricky to optimize. if a >= 0 this expression is false, but both parts of && play a role in concluding this.

I propose using ((a < 0) == (b > 0)) && (bool)(a%b)

• it has 6 operators, without Sign.
• for a common case where b is a known positive integer, this expression simply optimizes to (a < 0) && (bool)(a%b). If a is an expression that is not negative, this is false based on (a < 0) alone.

non-negative expressions are quite common (relu, argmax, sigmoid, ...), and this was found with a real model.