Add single-level TF2 APIs

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

@@ -65,6 +65,16 @@ public void testTf2() throws ClassHierarchyException, IllegalArgumentException,
 		testTf2("tf2e.py", "add", 2, 2, 3);
 		testTf2("tf2f.py", "add", 2, 2, 3);
 		testTf2("tf2g.py", "add", 2, 2, 3);
+		testTf2("tf2h.py", "add", 2, 2, 3);
+		testTf2("tf2i.py", "add", 2, 2, 3);
+		testTf2("tf2j.py", "add", 2, 2, 3);
+		testTf2("tf2k.py", "add", 2, 2, 3);
+		testTf2("tf2l.py", "add", 2, 2, 3);
+		testTf2("tf2m.py", "add", 2, 2, 3);
+		// TODO: Uncomment below test when https://github.com/wala/ML/issues/49 is fixed.
+		// testTf2("tf2n.py", "func2", 1, 2);
+		testTf2("tf2o.py", "add", 2, 2, 3);
+		testTf2("tf2p.py", "value_index", 2, 2, 3);
 	}
 
 	private void testTf2(String filename, String functionName, int expectedNumberOfTensorParameters, int... expectedValueNumbers) throws ClassHierarchyException, CancelException, IOException {

com.ibm.wala.cast.python.ml/source/com/ibm/wala/cast/python/ml/client/PythonTensorAnalysisEngine.java

@@ -112,7 +112,13 @@ private static Set<PointsToSetVariable> getDataflowSources(Graph<PointsToSetVari
 
 							// Single-level APIs.
 							if (Objects.equal(tensorFlowAPI, "ones") || Objects.equal(tensorFlowAPI, "Variable")
-									|| Objects.equal(tensorFlowAPI, "zeros") || Objects.equal(tensorFlowAPI, "constant"))
+									|| Objects.equal(tensorFlowAPI, "zeros") || Objects.equal(tensorFlowAPI, "constant")
+									|| Objects.equal(tensorFlowAPI, "SparseTensor") || Objects.equal(tensorFlowAPI, "Tensor")
+									|| Objects.equal(tensorFlowAPI, "fill") || Objects.equal(tensorFlowAPI, "eye")
+									|| Objects.equal(tensorFlowAPI, "zeros_like")
+									|| Objects.equal(tensorFlowAPI, "one_hot")
+									|| Objects.equal(tensorFlowAPI, "convert_to_tensor")
+									|| Objects.equal(tensorFlowAPI, "range"))
 								sources.add(src);
 							// Double-level APIs.
 							else if (Objects.equal(tensorFlowAPI, "uniform")) {

com.ibm.wala.cast.python.test/data/tf2h.py

@@ -0,0 +1,7 @@
+import tensorflow as tf
+
+def add(a, b):
+  return tf.sparse.add(a,b)
+
+
+c = add(tf.SparseTensor([[0, 0], [1, 2]], [1, 2],[3, 4]), tf.SparseTensor([[0, 0], [1, 2]], [1, 2], [3, 4]))

com.ibm.wala.cast.python.test/data/tf2i.py

@@ -0,0 +1,7 @@
+import tensorflow as tf
+
+def add(a, b):
+  return a + b
+
+
+c = add(tf.fill([1, 2], 2), tf.fill([2, 2], 1))

com.ibm.wala.cast.python.test/data/tf2j.py

@@ -0,0 +1,7 @@
+import tensorflow as tf
+
+def add(a, b):
+  return a + b
+
+
+c = add(tf.zeros_like([1, 2]), tf.zeros_like([2, 2]))

com.ibm.wala.cast.python.test/data/tf2k.py

@@ -0,0 +1,7 @@
+import tensorflow as tf
+
+def add(a, b):
+  return a + b
+
+
+c = add(tf.one_hot([0, 1, 2], 3), tf.one_hot([2,4,3], 3))

com.ibm.wala.cast.python.test/data/tf2l.py

@@ -0,0 +1,7 @@
+import tensorflow as tf
+
+def add(a, b):
+  return a + b
+
+
+c = add(tf.convert_to_tensor(1), tf.convert_to_tensor(2))

com.ibm.wala.cast.python.test/data/tf2m.py

@@ -0,0 +1,7 @@
+import tensorflow as tf
+
+def add(a, b):
+  return a + b
+
+
+c = add(tf.range(3, 18, 3), tf.range(5))

com.ibm.wala.cast.python.test/data/tf2n.py

@@ -0,0 +1,14 @@
+import tensorflow as tf
+
+def func2(t):
+  pass
+
+@tf.function
+def func():
+  a = tf.constant([[1.0, 2.0], [3.0, 4.0]])
+  b = tf.constant([[1.0, 1.0], [0.0, 1.0]])
+  c = tf.matmul(a, b)
+  tensor = tf.Tensor(c.op, 0, tf.float32)
+  func2(tensor)
+
+func()

com.ibm.wala.cast.python.test/data/tf2o.py

@@ -0,0 +1,7 @@
+import tensorflow as tf
+
+def add(a, b):
+  return tf.add(a,b)
+
+
+c = add(tf.eye(2,3), tf.eye(2,3))

com.ibm.wala.cast.python.test/data/tf2p.py

@@ -0,0 +1,13 @@
+import tensorflow as tf
+
+def value_index(a,b):
+  return a.value_index + b.value_index
+
+# From https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/Graph#using_graphs_directly_deprecated
+g = tf.Graph()
+with g.as_default():
+  # Defines operation and tensor in graph
+  c = tf.constant(30.0)
+  assert c.graph is g
+
+result = value_index(tf.Tensor(g.get_operations()[0], 0, tf.float32), tf.Tensor(g.get_operations()[0], 0, tf.float32))