Add a test case that uses datasets (#90)

Add test case for #89. We lose track of the tensors in the dataset.
wala · Oct 17, 2023 · 1b1ffac · 1b1ffac
1 parent 15a5341
commit 1b1ffac
Show file tree

Hide file tree

Showing 2 changed files with 17 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
@@ -195,6 +195,12 @@ public void testTf2()
     testTf2("tf2_testing_decorator8.py", "returned", 1, 3, 2);
     testTf2("tf2_testing_decorator9.py", "returned", 1, 3, 2);
     testTf2("tf2_testing_decorator10.py", "returned", 1, 3, 2);
+    testTf2(
+        "tf2_test_dataset.py",
+        "add",
+        0,
+        0); // NOTE: Change to testTf2("tf2_test_dataset.py", "add", 2, 3, 2, 3) once
+    // https://github.com/wala/ML/issues/89 is fixed.
   }
 
   private void testTf2(

diff --git a/com.ibm.wala.cast.python.test/data/tf2_test_dataset.py b/com.ibm.wala.cast.python.test/data/tf2_test_dataset.py
@@ -0,0 +1,11 @@
+import tensorflow as tf
+
+
+def add(a, b):
+  return a + b
+
+
+dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
+
+for element in dataset:
+    c = add(element, element)