Add dataset batch API.

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

@@ -206,6 +206,7 @@ public void testTf2()
     testTf2("tf2_test_dataset2.py", "add", 2, 2, 2, 3);
     testTf2("tf2_test_dataset3.py", "add", 2, 2, 2, 3);
     testTf2("tf2_test_dataset4.py", "add", 2, 2, 2, 3);
+    testTf2("tf2_test_dataset5.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);

com.ibm.wala.cast.python.ml/data/tensorflow.xml

@@ -796,6 +796,8 @@
         <method name="read_dataset" descriptor="()LRoot;">
           <new def="shuffle" class="Ltensorflow/data/shuffle" />
           <putfield class="LRoot" field="shuffle" fieldType="LRoot" ref="arg0" value="shuffle" />
+          <new def="batch" class="Ltensorflow/data/batch" />
+          <putfield class="LRoot" field="batch" fieldType="LRoot" ref="arg0" value="batch" />
           <return value="arg0" />
         </method>
         <method name="do" descriptor="()LRoot;" numArgs="2" paramNames="self variant_tensor">
@@ -813,6 +815,16 @@
           <return value="xx" />
         </method>
       </class>
+
+      <class name="batch" allocatable="true">
+        <!-- https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/data/Dataset#batch-->
+        <method name="do" descriptor="()LRoot;" numArgs="6" paramNames="self batch_size drop_remainder num_parallel_calls deterministic name">
+          <!-- FIXME: Workaround for https://github.com/wala/ML/issues/127. -->
+          <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/data/Dataset">

com.ibm.wala.cast.python.test/data/tf2_test_dataset5.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]).shuffle(3).batch(2)
+
+for element in dataset:
+    c = add(element, element)