Add test for varying dtypes.

Author: khatchad

com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java

@@ -291,6 +291,17 @@ public void testDecorator11()
     test("tf2_test_decorator11.py", "C.returned", 1, 1, Map.of(3, Set.of(TENSOR_5_INT32)));
   }
 
+  @Test
+  public void testDecorator12()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test(
+        "tf2_test_decorator12.py",
+        "returned",
+        1,
+        1,
+        Map.of(2, Set.of(TENSOR_2_FLOAT32, TENSOR_2_INT32)));
+  }
+
   @Test
   public void testDataset()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {

com.ibm.wala.cast.python.ml/source/com/ibm/wala/cast/python/ml/client/Range.java

@@ -130,6 +130,10 @@ protected EnumSet<DType> getDefaultDTypes(PropagationCallGraphBuilder builder) {
             .distinct()
             .collect(Collectors.toCollection(() -> EnumSet.noneOf(DType.class)));
 
+    // FIXME: We can't tell the difference here between varying dtypes in a single call and that of
+    // possible varying dtypes values from the points-to graph. Below, we are treating it as these
+    // values lie in a single call, but that may not be the case.
+
     if (types.contains(DType.FLOAT64)) return EnumSet.of(DType.FLOAT64);
     else if (types.contains(DType.FLOAT32)) return EnumSet.of(DType.FLOAT32);
     else if (types.contains(DType.INT64)) return EnumSet.of(DType.INT64);

com.ibm.wala.cast.python.test/data/tf2_test_decorator12.py

@@ -0,0 +1,14 @@
+import tensorflow as tf
+
+
+@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.float32),))
+@tf.function(reduce_retracing=True)
+def returned(a):
+    return a
+
+
+a = tf.constant([1, 1.0])
+b = returned(a)
+
+assert a.shape == (2,)
+assert a.dtype == tf.float32