[Feature] Add support for tf.math.xlogy

[jaytiwarihub]

Description:
This PR adds support for the tf.math.xlogy(x, y) operation, which computes x * log(y) but returns 0 if x == 0. This is a common operation in loss functions (like cross-entropy) and probabilistic models.

Implementation Details:

• Mapped Xlogy to a combination of ONNX nodes: Log, Mul, and Where.
• Implemented safe handling for the x=0 case:
  • Standard x * log(y) results in NaN if x=0 and y<0.
  • Used Where(x==0, 0, x * log(y)) to ensure the result is correctly 0 (matching TensorFlow behavior) and to avoid propagating NaNs.

Changes:

• Added Xlogy handler to tf2onnx/onnx_opset/math.py.

Verification:

• Verified locally with a custom script comparing tf.math.xlogy output against the converted ONNX model.
• Confirmed correct graph generation for standard inputs and the x=0 edge case.

[sonarqubecloud]

Quality Gate passed

Issues
0 New issues
0 Accepted issues

Measures
0 Security Hotspots
0.0% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud

[Copilot]

Pull request overview

Adds a TensorFlow op handler for tf.math.xlogy(x, y) to the ONNX conversion pipeline, aiming to preserve TensorFlow’s x==0 => 0 semantics (avoiding 0 * log(y) NaNs in edge cases).

Changes:

• Introduces a new @tf_op("Xlogy") handler in tf2onnx/onnx_opset/math.py.
• Builds an ONNX subgraph using Log, Mul, Equal, and Where to implement xlogy.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The inserted Log node won’t go through the existing DirectOp.version_6 workaround that casts DOUBLE inputs to FLOAT because some runtimes don’t implement Log on doubles. This means Xlogy may regress for float64 inputs even though tf.math.log is handled elsewhere. Consider applying the same cast-to-float/log/cast-back logic around log(y) when y is DOUBLE (and keep Mul inputs’ dtypes consistent).

Suggested change

	# Compute x * log(y)
	log_y = ctx.make_node("Log", [y])
	x_log_y = ctx.make_node("Mul", [x, log_y.output[0]])
	# Create zero constant matching x's dtype
	dtype = ctx.get_dtype(x)
	zero_val = np.array(0, dtype=utils.map_onnx_to_numpy_type(dtype))
	zero_const = ctx.make_const(utils.make_name("zero"), zero_val)
	# If x == 0, return 0. Otherwise return x * log(y).
	condition = ctx.make_node("Equal", [x, zero_const.output[0]])
	node = ctx.make_node("Where",
	[condition.output[0], zero_const.output[0], x_log_y.output[0]],
	name=node.name)
	# Compute x * log(y)
	y_dtype = ctx.get_dtype(y)
	if y_dtype == onnx_pb.TensorProto.DOUBLE:
	# Some runtimes don't implement Log on DOUBLE: cast to FLOAT, apply Log, cast back.
	y_cast = ctx.make_node(
	"Cast",
	[y],
	attr={"to": onnx_pb.TensorProto.FLOAT},
	)
	ctx.set_dtype(y_cast.output[0], onnx_pb.TensorProto.FLOAT)
	log_y_inner = ctx.make_node("Log", [y_cast.output[0]])
	log_y = ctx.make_node(
	"Cast",
	[log_y_inner.output[0]],
	attr={"to": y_dtype},
	)
	ctx.set_dtype(log_y.output[0], y_dtype)
	else:
	log_y = ctx.make_node("Log", [y])
	x_log_y = ctx.make_node("Mul", [x, log_y.output[0]])
	# Create zero constant matching x's dtype
	dtype = ctx.get_dtype(x)
	zero_val = np.array(0, dtype=utils.map_onnx_to_numpy_type(dtype))
	zero_const = ctx.make_const(utils.make_name("zero"), zero_val)
	# If x == 0, return 0. Otherwise return x * log(y).
	condition = ctx.make_node("Equal", [x, zero_const.output[0]])
	node = ctx.make_node(
	"Where",
	[condition.output[0], zero_const.output[0], x_log_y.output[0]],
	name=node.name,
	)

[Copilot]

This adds a new op mapping but there doesn’t appear to be any automated test coverage for tf.math.xlogy yet. Adding a backend test (similar to test_log_double / test_div_no_nan) would help lock in behavior for the x==0 edge case (including y<0) and for float64 inputs.

Suggested change

	# Compute x * log(y)
	log_y = ctx.make_node("Log", [y])
	x_log_y = ctx.make_node("Mul", [x, log_y.output[0]])
	# Create zero constant matching x's dtype
	dtype = ctx.get_dtype(x)
	zero_val = np.array(0, dtype=utils.map_onnx_to_numpy_type(dtype))
	zero_const = ctx.make_const(utils.make_name("zero"), zero_val)
	# If x == 0, return 0. Otherwise return x * log(y).
	condition = ctx.make_node("Equal", [x, zero_const.output[0]])
	node = ctx.make_node("Where",
	[condition.output[0], zero_const.output[0], x_log_y.output[0]],
	name=node.name)
	return node.output
	# Determine dtype of x (used as output dtype)
	dtype_x = ctx.get_dtype(x)
	# For double tensors, avoid using Log on DOUBLE by computing in FLOAT
	if dtype_x == onnx_pb.TensorProto.DOUBLE:
	# Cast inputs to FLOAT for the computation
	x_cast = ctx.make_node("Cast", [x], to=onnx_pb.TensorProto.FLOAT)
	y_cast = ctx.make_node("Cast", [y], to=onnx_pb.TensorProto.FLOAT)
	# Compute x * log(y) in FLOAT
	log_y = ctx.make_node("Log", [y_cast.output[0]])
	x_log_y = ctx.make_node("Mul", [x_cast.output[0], log_y.output[0]])
	# Create zero constant in FLOAT
	zero_val = np.array(0, dtype=np.float32)
	zero_const = ctx.make_const(utils.make_name("zero"), zero_val)
	# If x == 0, return 0. Otherwise return x * log(y), all in FLOAT
	condition = ctx.make_node("Equal", [x_cast.output[0], zero_const.output[0]])
	where_node = ctx.make_node(
	"Where",
	[condition.output[0], zero_const.output[0], x_log_y.output[0]],
	name=node.name + "_float"
	)
	# Cast result back to DOUBLE to match original dtype
	cast_back_node = ctx.make_node(
	"Cast",
	[where_node.output[0]],
	to=onnx_pb.TensorProto.DOUBLE,
	name=node.name
	)
	return cast_back_node.output
	# Non-DOUBLE path: compute directly in the original dtype
	# Compute x * log(y)
	log_y = ctx.make_node("Log", [y])
	x_log_y = ctx.make_node("Mul", [x, log_y.output[0]])
	# Create zero constant matching x's dtype
	zero_val = np.array(0, dtype=utils.map_onnx_to_numpy_type(dtype_x))
	zero_const = ctx.make_const(utils.make_name("zero"), zero_val)
	# If x == 0, return 0. Otherwise return x * log(y).
	condition = ctx.make_node("Equal", [x, zero_const.output[0]])
	where_node = ctx.make_node(
	"Where",
	[condition.output[0], zero_const.output[0], x_log_y.output[0]],
	name=node.name
	)
	return where_node.output

[Copilot]

This handler creates new ONNX nodes but never rewires the original TF Xlogy node’s output to them (and handler return values are ignored by tensorflow_onnx_mapping). As a result, the graph will still contain an unconverted Xlogy node and downstream ops will continue to consume its original output. Convert in-place (e.g., change node to Mul, replace its second input with Log(y), then insert a Where on node.output[0] like DivNoNan does) or remove the original node and recreate the final node with outputs=node.output.

Suggested change

	# Compute x * log(y)
	log_y = ctx.make_node("Log", [y])
	x_log_y = ctx.make_node("Mul", [x, log_y.output[0]])
	# Create zero constant matching x's dtype
	dtype = ctx.get_dtype(x)
	zero_val = np.array(0, dtype=utils.map_onnx_to_numpy_type(dtype))
	zero_const = ctx.make_const(utils.make_name("zero"), zero_val)
	# If x == 0, return 0. Otherwise return x * log(y).
	condition = ctx.make_node("Equal", [x, zero_const.output[0]])
	node = ctx.make_node("Where",
	[condition.output[0], zero_const.output[0], x_log_y.output[0]],
	name=node.name)
	return node.output
	# Compute x * log(y)
	log_y = ctx.make_node("Log", [y])
	x_log_y = ctx.make_node("Mul", [x, log_y.output[0]])
	# Create zero constant matching x's dtype
	dtype = ctx.get_dtype(x)
	zero_val = np.array(0, dtype=utils.map_onnx_to_numpy_type(dtype))
	zero_const = ctx.make_const(utils.make_name("zero"), zero_val)
	# If x == 0, return 0. Otherwise return x * log(y).
	condition = ctx.make_node("Equal", [x, zero_const.output[0]])
	# Replace the original Xlogy node by a Where node that produces the same outputs.
	where_node = ctx.make_node(
	"Where",
	[condition.output[0], zero_const.output[0], x_log_y.output[0]],
	name=node.name,
	outputs=node.output,
	)
	# Remove the original Xlogy node from the graph.
	ctx.remove_node(node.name)

[Copilot]

Where is only available starting with ONNX opset 9, but this handler is registered as version_1. That will generate invalid models when targeting opsets < 9 (and makes it unclear which minimum opset is required). Register this as version_9 (and optionally add a lower-opset fallback or an explicit failure path).

Suggested change

	x = node.input[0]
	y = node.input[1]
	# Compute x * log(y)
	log_y = ctx.make_node("Log", [y])
	x_log_y = ctx.make_node("Mul", [x, log_y.output[0]])
	# Create zero constant matching x's dtype
	dtype = ctx.get_dtype(x)
	zero_val = np.array(0, dtype=utils.map_onnx_to_numpy_type(dtype))
	zero_const = ctx.make_const(utils.make_name("zero"), zero_val)
	# If x == 0, return 0. Otherwise return x * log(y).
	condition = ctx.make_node("Equal", [x, zero_const.output[0]])
	node = ctx.make_node("Where",
	[condition.output[0], zero_const.output[0], x_log_y.output[0]],
	name=node.name)
	# Xlogy uses the ONNX Where operator, which is only available starting with opset 9.
	# For opsets < 9, explicitly fail instead of generating an invalid model.
	utils.make_sure(False, "Xlogy is only supported for ONNX opset >= 9")
	@classmethod
	def version_9(cls, ctx, node, **kwargs):
	x = node.input[0]
	y = node.input[1]
	# Compute x * log(y)
	log_y = ctx.make_node("Log", [y])
	x_log_y = ctx.make_node("Mul", [x, log_y.output[0]])
	# Create zero constant matching x's dtype
	dtype = ctx.get_dtype(x)
	zero_val = np.array(0, dtype=utils.map_onnx_to_numpy_type(dtype))
	zero_const = ctx.make_const(utils.make_name("zero"), zero_val)
	# If x == 0, return 0. Otherwise return x * log(y).
	condition = ctx.make_node("Equal", [x, zero_const.output[0]])
	node = ctx.make_node(
	"Where",
	[condition.output[0], zero_const.output[0], x_log_y.output[0]],
	name=node.name,
	)

[Copilot]

This uses ctx.get_dtype(x) and immediately passes it to utils.map_onnx_to_numpy_type. If the dtype is unknown at conversion time (None), this will throw. Other code in this file guards this case with utils.make_sure(dtype is not None, ...); consider adding the same guard (or deriving the dtype from node.output_dtypes[0] if available).

Suggested change

	dtype = ctx.get_dtype(x)
	dtype = ctx.get_dtype(x)
	if dtype is None and getattr(node, "output_dtypes", None):
	# Fall back to the node's output dtype if available
	dtype = node.output_dtypes[0]
	utils.make_sure(dtype is not None, "Xlogy: cannot determine dtype for input '%s'", x)

[jaytiwarihub]

@xadupre I've really enjoyed diving into the graph rewriting and adding missing ops over the last few weeks. I was wondering: is tensorflow-onnx planning to participate in the upcoming LFX Summer Mentorship term? If so, I would love to officially apply and take on a larger feature.