From 78a9de9c22b06da2724e2a6ba13f2493025058d6 Mon Sep 17 00:00:00 2001
From: Raffi Khatchadourian <khatchad@hunter.cuny.edu>
Date: Wed, 31 Jan 2024 10:51:07 -0500
Subject: [PATCH] New API and new operator overloading issue (#81)

* Add `tf.nn.sparse_softmax_cross_entropy_with_logits()`.

* Mark test that tests https://github.com/wala/ML/issues/136.

* Update tests to reflect new API.

* Update test in light of https://github.com/wala/ML/issues/136.
---
 .../cast/python/ml/test/TestTensorflowModel.java | 16 +++++++++++++---
 com.ibm.wala.cast.python.ml/data/tensorflow.xml  |  8 ++++++++
 ...t_sparse_softmax_cross_entropy_with_logits.py | 14 ++++++++++++++
 .../data/tf2_test_tensor_list3.py                |  2 ++
 4 files changed, 37 insertions(+), 3 deletions(-)
 create mode 100644 com.ibm.wala.cast.python.test/data/tf2_test_sparse_softmax_cross_entropy_with_logits.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 ad7ed7c16..50fab18f9 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
@@ -217,7 +217,11 @@ public void testTf2()
     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);
+    testTf2(
+        "tf2_test_tensor_list3.py",
+        "add",
+        0,
+        0); // NOTE: Change to 2, 2, 2, 3 once https://github.com/wala/ML/issues/136 is fixed.
     testTf2("tf2_test_tensor_list4.py", "add", 0, 0);
     testTf2("tf2_test_tensor_list5.py", "add", 0, 0);
     testTf2("tf2_test_model_call.py", "SequentialModel.__call__", 1, 1, 3);
@@ -234,10 +238,10 @@ public void testTf2()
         "neural_network.py",
         "cross_entropy_loss",
         1,
-        4,
+        8,
         3); // NOTE: Change to 2 tensor parameters once https://github.com/wala/ML/issues/127 is
     // fixed. Values 2 and 3 will correspond to the tensor parameters.
-    testTf2("neural_network.py", "run_optimization", 2, 2, 2, 3);
+    testTf2("neural_network.py", "run_optimization", 2, 3, 2, 3);
     testTf2(
         "neural_network.py",
         "accuracy",
@@ -259,6 +263,11 @@ public void testTf2()
     testTf2("tf2_test_add5.py", "f", 1, 1, 2);
     testTf2("tf2_test_add6.py", "f", 1, 1, 2);
     testTf2("multigpu_training.py", "run_optimization", 2, 4, 2, 3);
+    testTf2(
+        "multigpu_training.py",
+        "average_gradients",
+        0,
+        0); // NOTE: Change to 1, 1, 2 once https://github.com/wala/ML/issues/136 is fixed.
     testTf2("tf2_test_reduce_mean.py", "f", 1, 1, 2);
     testTf2("tf2_test_reduce_mean.py", "g", 1, 1, 2);
     testTf2("tf2_test_reduce_mean.py", "h", 1, 1, 2);
@@ -266,6 +275,7 @@ public void testTf2()
     testTf2("tf2_test_gradient2.py", "f", 1, 1, 2);
     testTf2("tf2_test_multiply.py", "f", 1, 1, 2);
     testTf2("tf2_test_multiply2.py", "f", 1, 1, 2);
+    testTf2("tf2_test_sparse_softmax_cross_entropy_with_logits.py", "f", 1, 1, 2);
   }
 
   private void testTf2(
diff --git a/com.ibm.wala.cast.python.ml/data/tensorflow.xml b/com.ibm.wala.cast.python.ml/data/tensorflow.xml
index c6d532d53..224995705 100644
--- a/com.ibm.wala.cast.python.ml/data/tensorflow.xml
+++ b/com.ibm.wala.cast.python.ml/data/tensorflow.xml
@@ -75,6 +75,8 @@
         <putfield class="LRoot" field="conv3d" fieldType="LRoot" ref="nn" value="conv3d" />
         <new def="softmax" class="Ltensorflow/functions/softmax" />
         <putfield class="LRoot" field="softmax" fieldType="LRoot" ref="nn" value="softmax" />
+        <new def="sparse_softmax_cross_entropy_with_logits" class="Ltensorflow/functions/sparse_softmax_cross_entropy_with_logits" />
+        <putfield class="LRoot" field="sparse_softmax_cross_entropy_with_logits" fieldType="LRoot" ref="nn" value="sparse_softmax_cross_entropy_with_logits" />
         <new def="sigmoid" class="Ltensorflow/math/sigmoid" />
         <putfield class="LRoot" field="sigmoid" fieldType="LRoot" ref="nn" value="sigmoid" />
         <putfield class="LRoot" field="sigmoid" fieldType="LRoot" ref="math" value="sigmoid" />
@@ -697,6 +699,12 @@
           <return value="logits" />
         </method>
       </class>
+      <class name="sparse_softmax_cross_entropy_with_logits" allocatable="true">
+        <!-- https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/nn/sparse_softmax_cross_entropy_with_logits -->
+        <method name="do" descriptor="()LRoot;" numArgs="4" paramNames="self labels logits name">
+          <return value="labels" />
+        </method>
+      </class>
     </package>
     <package name="tensorflow/estimator">
       <class name="Estimator" allocatable="true">
diff --git a/com.ibm.wala.cast.python.test/data/tf2_test_sparse_softmax_cross_entropy_with_logits.py b/com.ibm.wala.cast.python.test/data/tf2_test_sparse_softmax_cross_entropy_with_logits.py
new file mode 100644
index 000000000..86a0dd604
--- /dev/null
+++ b/com.ibm.wala.cast.python.test/data/tf2_test_sparse_softmax_cross_entropy_with_logits.py
@@ -0,0 +1,14 @@
+# from https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/nn/sparse_softmax_cross_entropy_with_logits
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+logits = tf.constant(
+    [[2.0, -5.0, 0.5, -0.1], [0.0, 0.0, 1.9, 1.4], [-100.0, 100.0, -100.0, -100.0]]
+)
+labels = tf.constant([0, 3, 1])
+f(tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=logits.numpy()))
diff --git a/com.ibm.wala.cast.python.test/data/tf2_test_tensor_list3.py b/com.ibm.wala.cast.python.test/data/tf2_test_tensor_list3.py
index fb674b2e0..d00832a2f 100644
--- a/com.ibm.wala.cast.python.test/data/tf2_test_tensor_list3.py
+++ b/com.ibm.wala.cast.python.test/data/tf2_test_tensor_list3.py
@@ -1,3 +1,5 @@
+# Test https://github.com/wala/ML/issues/136.
+
 import tensorflow as tf