From 2f42a051efad28b97d06cfbffcc5bf08da1420b8 Mon Sep 17 00:00:00 2001
From: Raffi Khatchadourian <khatchad@hunter.cuny.edu>
Date: Mon, 29 Jun 2026 15:07:53 -0400
Subject: [PATCH] Frame g's parameter as the multi-element set in the
 input_signature probe (wala/ML#638).

Follow-up to the merged #467. The decorated `f` passes its parameter to `g`; what
`g` receives depends on the execution mode a static analysis cannot determine:
traced the signature governs ((None,) int32), eager it is the call-site argument
((3,) int32). So the sound type of `g`'s parameter is the set {(None,), (3,)}.
The fixture now documents both modes without a `run_functions_eagerly` call (the
eager element is the argument at the call site), and the test pins the current
argument-derived (3,) with the sound set spelled out in the TODO toward
wala/ML#638.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01USFp76oHfj2u3wUSCo87cg
---
 .../cast/python/ml/test/TestTensorflow2Model.java  | 14 +++++++++-----
 .../tf2_test_decorated_call_depth_input_sig.py     | 11 ++++++++---
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java b/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java
index 926d1c5b9..7a3803152 100644
--- a/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java
+++ b/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java
@@ -4738,11 +4738,15 @@ public void testDecoratedCallDepth()
 
   /**
    * A {@code @tf.function(input_signature=[(None,) int32])} passes its parameter to {@code g} (the
-   * FUT). At runtime {@code g} receives the signature's {@code (None,)} int32, since the signature
-   * governs the parameter and propagates to the callee.
-   *
-   * <p>TODO: Ariadne ignores {@code input_signature} and types {@code g}'s parameter from the
-   * call-site argument {@code (3,)} int32, which is <em>unsound</em> here. Tracked by <a
+   * FUT). What {@code g} receives depends on the execution mode, which a static analysis cannot
+   * determine: traced (the default) the signature governs and {@code g} receives {@code (None,)}
+   * int32; under {@code run_functions_eagerly} the signature is ignored and {@code g} receives the
+   * call-site argument's {@code (3,)} int32. So the sound type of {@code g}'s parameter is the set
+   * {@code {(None,), (3,)}} int32.
+   *
+   * <p>TODO: this pins the current behavior. Ariadne does not consume {@code input_signature}, so
+   * it produces only the argument-derived {@code (3,)} element and misses the signature-derived
+   * {@code (None,)} one; the sound result is the set {@code {(None,), (3,)}} int32. Tracked by <a
    * href="https://github.com/wala/ML/issues/638">wala/ML#638</a>.
    */
   @Test
diff --git a/com.ibm.wala.cast.python.test/data/tf2_test_decorated_call_depth_input_sig.py b/com.ibm.wala.cast.python.test/data/tf2_test_decorated_call_depth_input_sig.py
index 3e9d4605f..885467431 100644
--- a/com.ibm.wala.cast.python.test/data/tf2_test_decorated_call_depth_input_sig.py
+++ b/com.ibm.wala.cast.python.test/data/tf2_test_decorated_call_depth_input_sig.py
@@ -7,11 +7,16 @@ def g(b):
 
 @tf.function(input_signature=[tf.TensorSpec(shape=(None,), dtype=tf.int32)])
 def f(x):
-    # `input_signature` governs `x`, so at `g`'s call site the argument is (None,) int32,
-    # NOT the (3,) of the value passed to `f` below.
+    # Traced (the default): `input_signature` governs `x`, so `g` receives (None,) int32.
     assert x.shape.as_list() == [None]
     assert x.dtype == tf.int32
     g(x)
 
 
-f(tf.constant([1, 2, 3], dtype=tf.int32))
+# Under `run_functions_eagerly` the signature would be ignored and `g` would instead receive this
+# argument's (3,) int32. A static analysis cannot know the execution mode, so the sound type of
+# `g`'s parameter is the set {(None,), (3,)} int32.
+arg = tf.constant([1, 2, 3], dtype=tf.int32)
+assert arg.shape == (3,)
+assert arg.dtype == tf.int32
+f(arg)