From 0a8487213b4a5292e2685bd28ab01bc5e55018e8 Mon Sep 17 00:00:00 2001
From: Raffi Khatchadourian <khatchad@hunter.cuny.edu>
Date: Fri, 26 Jan 2024 16:11:30 -0500
Subject: [PATCH] Add dataset generator and test.

---
 .../python/ml/test/TestTensorflowModel.java   |  1 +
 .../data/tensorflow.xml                       | 10 ++++++++++
 .../data/tf2_test_dataset10.py                | 20 +++++++++++++++++++
 3 files changed, 31 insertions(+)
 create mode 100644 com.ibm.wala.cast.python.test/data/tf2_test_dataset10.py

diff --git a/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflowModel.java b/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflowModel.java
index f42c8a7da..eafe1a115 100644
--- a/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflowModel.java
+++ b/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflowModel.java
@@ -214,6 +214,7 @@ public void testTf2()
     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_dataset10.py", "add", 2, 2, 2, 3);
     testTf2("tf2_test_tensor_list.py", "add", 2, 2, 2, 3);
     testTf2("tf2_test_tensor_list2.py", "add", 0, 0);
     testTf2("tf2_test_tensor_list3.py", "add", 0, 0);
diff --git a/com.ibm.wala.cast.python.ml/data/tensorflow.xml b/com.ibm.wala.cast.python.ml/data/tensorflow.xml
index 60a30dd4d..8cc2c550c 100644
--- a/com.ibm.wala.cast.python.ml/data/tensorflow.xml
+++ b/com.ibm.wala.cast.python.ml/data/tensorflow.xml
@@ -61,6 +61,8 @@
         <putfield class="LRoot" field="numpy_input_fn" fieldType="LRoot" ref="inputs" value="numpy_input_fn" />
         <new def="from_tensor_slices" class="Ltensorflow/data/Dataset/from_tensor_slices" />
         <putfield class="LRoot" field="from_tensor_slices" fieldType="LRoot" ref="Dataset" value="from_tensor_slices" />
+        <new def="from_generator" class="Ltensorflow/data/Dataset/from_generator" />
+        <putfield class="LRoot" field="from_generator" fieldType="LRoot" ref="Dataset" value="from_generator" />
         <new def="reshape" class="Ltensorflow/functions/reshape" />
         <putfield class="LRoot" field="reshape" fieldType="LRoot" ref="x" value="reshape" />
         <new def="conv2d" class="Ltensorflow/functions/conv2d" />
@@ -790,6 +792,14 @@
           <return value="xx" />
         </method>
       </class>
+      <class name="from_generator" allocatable="true">
+        <!-- https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/data/Dataset#from_generator -->
+        <method name="do" descriptor="()LRoot;" numArgs="6" paramNames="generator output_types output_shapes args output_signature name">
+          <new def="x" class="Ltensorflow/data/Dataset" />
+          <call class="Ltensorflow/data/Dataset" name="read_dataset" descriptor="()LRoot;" type="virtual" arg0="x" def="xx" />
+          <return value="xx" />
+        </method>
+      </class>
     </package>
     <package name="tensorflow/estimator/train">
       <class name="train" allocatable="true">
diff --git a/com.ibm.wala.cast.python.test/data/tf2_test_dataset10.py b/com.ibm.wala.cast.python.test/data/tf2_test_dataset10.py
new file mode 100644
index 000000000..bf7392bfb
--- /dev/null
+++ b/com.ibm.wala.cast.python.test/data/tf2_test_dataset10.py
@@ -0,0 +1,20 @@
+import tensorflow as tf
+
+
+def gen():
+  ragged_tensor = tf.ragged.constant([[1, 2], [3]])
+  yield 42, ragged_tensor
+
+
+def add(a, b):
+  return a + b
+
+
+dataset = tf.data.Dataset.from_generator(
+     gen,
+     output_signature=(
+         tf.TensorSpec(shape=(), dtype=tf.int32),
+         tf.RaggedTensorSpec(shape=(2, None), dtype=tf.int32)))
+
+for element in dataset:
+    c = add(element, element)