Progress.

Author: khatchad

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

@@ -1791,13 +1791,23 @@ public void testAdd108()
   @Test
   public void testAdd109()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
-    test("tf2_test_add109.py", "add", 2, 2, Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
+    test(
+        "tf2_test_add109.py",
+        "add",
+        2,
+        2,
+        Map.of(2, Set.of(TENSOR_2_FLOAT32), 3, Set.of(TENSOR_2_FLOAT32)));
   }
 
   @Test
   public void testAdd110()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
-    test("tf2_test_add110.py", "add", 2, 2, Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
+    test(
+        "tf2_test_add110.py",
+        "add",
+        2,
+        2,
+        Map.of(2, Set.of(TENSOR_2_FLOAT32), 3, Set.of(TENSOR_2_FLOAT32)));
   }
 
   @Test

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

@@ -5,6 +5,7 @@
 import static com.ibm.wala.cast.python.ml.types.TensorFlowTypes.NORMAL;
 import static com.ibm.wala.cast.python.ml.types.TensorFlowTypes.ONES;
 import static com.ibm.wala.cast.python.ml.types.TensorFlowTypes.RANGE;
+import static com.ibm.wala.cast.python.ml.types.TensorFlowTypes.TRUNCATED_NORMAL;
 import static com.ibm.wala.cast.python.ml.types.TensorFlowTypes.UNIFORM;
 import static com.ibm.wala.cast.python.ml.types.TensorFlowTypes.ZEROS;
 import static com.ibm.wala.cast.python.ml.types.TensorFlowTypes.ZEROS_LIKE;
@@ -40,6 +41,8 @@ public static TensorGenerator getGenerator(PointsToSetVariable source) {
     else if (calledFunction.equals(RANGE.getDeclaringClass())) return new Range(source, node);
     else if (calledFunction.equals(UNIFORM.getDeclaringClass())) return new Uniform(source, node);
     else if (calledFunction.equals(NORMAL.getDeclaringClass())) return new Normal(source, node);
+    else if (calledFunction.equals(TRUNCATED_NORMAL.getDeclaringClass()))
+      return new TruncatedNormal(source, node);
     else if (calledFunction.equals(ZEROS.getDeclaringClass())) return new Zeros(source, node);
     else if (calledFunction.equals(ZEROS_LIKE.getDeclaringClass()))
       return new ZerosLike(source, node);

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

@@ -0,0 +1,20 @@
+package com.ibm.wala.cast.python.ml.client;
+
+import com.ibm.wala.ipa.callgraph.CGNode;
+import com.ibm.wala.ipa.callgraph.propagation.PointsToSetVariable;
+
+/**
+ * A representation of the `tf.random.truncated_normal' API in TensorFlow.
+ *
+ * @see <a
+ *     href="https://www.tensorflow.org/api_docs/python/tf/random/truncated_normal">tf.random.truncated_normal
+ *     API</a>.
+ * @author <a href="mailto:[email protected]">Raffi Khatchadourian</a>
+ */
+public class TruncatedNormal extends Normal {
+
+  public TruncatedNormal(PointsToSetVariable source, CGNode node) {
+    super(source, node);
+    // TODO Auto-generated constructor stub
+  }
+}

com.ibm.wala.cast.python.ml/source/com/ibm/wala/cast/python/ml/types/TensorFlowTypes.java

@@ -83,6 +83,14 @@ public enum DType {
               PythonTypes.pythonLoader, TypeName.string2TypeName("Ltensorflow/functions/normal")),
           AstMethodReference.fnSelector);
 
+  /** https://www.tensorflow.org/api_docs/python/tf/random/truncated_normal. */
+  public static final MethodReference TRUNCATED_NORMAL =
+      MethodReference.findOrCreate(
+          TypeReference.findOrCreate(
+              PythonTypes.pythonLoader,
+              TypeName.string2TypeName("Ltensorflow/functions/truncated_normal")),
+          AstMethodReference.fnSelector);
+
   /** https://www.tensorflow.org/api_docs/python/tf/zeros. */
   public static final MethodReference ZEROS =
       MethodReference.findOrCreate(

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

@@ -4,5 +4,9 @@
 def add(a, b):
     return a + b
 
+arg1 = tf.random.truncated_normal([2])
+assert isinstance(arg1, tf.Tensor)
+assert arg1.dtype == tf.float32
+assert arg1.shape == (2,)
 
-c = add(tf.random.truncated_normal([2]), tf.random.truncated_normal([2], 3, 1))
+c = add(arg1, tf.random.truncated_normal([2], 3, 1))