Add interprocedural test case.

ponder-lab · Jan 5, 2024 · fcbfa42 · fcbfa42
1 parent e9e3f57
commit fcbfa42
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/...wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflowModel.java b/...wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflowModel.java
@@ -210,6 +210,7 @@ public void testTf2()
     testTf2("tf2_test_dataset6.py", "add", 2, 2, 2, 3);
     testTf2("tf2_test_dataset7.py", "add", 2, 2, 2, 3);
     testTf2("tf2_test_dataset8.py", "add", 2, 2, 2, 3);
+    testTf2("tf2_test_dataset9.py", "add", 2, 2, 2, 3);
     testTf2("tf2_test_tensor_list.py", "add", 2, 3, 2, 3);
     testTf2("tf2_test_tensor_list2.py", "add", 0, 2);
     testTf2("tf2_test_tensor_list3.py", "add", 0, 2);

diff --git a/com.ibm.wala.cast.python.test/data/tf2_test_dataset9.py b/com.ibm.wala.cast.python.test/data/tf2_test_dataset9.py
@@ -0,0 +1,15 @@
+import tensorflow as tf
+
+
+def add(a, b):
+    return a + a
+
+
+def func(ds):
+    for images, labels in ds:
+        c = add(images, labels)
+
+
+(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
+dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000).batch(32)
+func(dataset)