Handle tf.random.poisson API in TensorFlow.

Author: khatchad

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

@@ -1884,13 +1884,23 @@ public void testAdd106()
   @Test
   public void testAdd107()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
-    test("tf2_test_add107.py", "add", 2, 2, Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
+    test(
+        "tf2_test_add107.py",
+        "add",
+        2,
+        2,
+        Map.of(2, Set.of(TENSOR_10_2_FLOAT32), 3, Set.of(TENSOR_10_2_FLOAT32)));
   }
 
   @Test
   public void testAdd108()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
-    test("tf2_test_add108.py", "add", 2, 2, Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
+    test(
+        "tf2_test_add108.py",
+        "add",
+        2,
+        2,
+        Map.of(2, Set.of(TENSOR_10_2_FLOAT32), 3, Set.of(TENSOR_10_2_FLOAT32)));
   }
 
   @Test
@@ -4471,6 +4481,30 @@ public void testGamma6()
     test("tf2_test_gamma6.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_30_3_2_FLOAT32)));
   }
 
+  @Test
+  public void testPoisson()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_poisson.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_10_2_FLOAT32)));
+  }
+
+  @Test
+  public void testPoisson2()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_poisson2.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_10_2_FLOAT32)));
+  }
+
+  @Test
+  public void testPoisson3()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_poisson3.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_10_2_FLOAT64)));
+  }
+
+  @Test
+  public void testPoisson4()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_poisson4.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_7_5_2_FLOAT32)));
+  }
+
   private void test(
       String filename,
       String functionName,

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

@@ -0,0 +1,90 @@
+package com.ibm.wala.cast.python.ml.client;
+
+import static com.ibm.wala.cast.python.ml.client.Poisson.Parameters.DTYPE;
+import static com.ibm.wala.cast.python.ml.client.Poisson.Parameters.LAM;
+
+import com.ibm.wala.cast.python.ml.types.TensorType.Dimension;
+import com.ibm.wala.ipa.callgraph.CGNode;
+import com.ibm.wala.ipa.callgraph.propagation.PointsToSetVariable;
+import com.ibm.wala.ipa.callgraph.propagation.PropagationCallGraphBuilder;
+import com.ibm.wala.util.collections.HashSetFactory;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+
+/**
+ * A representation of the `tf.random.poisson` API in TensorFlow.
+ *
+ * @see <a
+ *     href="https://www.tensorflow.org/api_docs/python/tf/random/poisson">tf.random.poisson</a>.
+ * @author <a href="mailto:[email protected]">Raffi Khatchadourian</a>
+ */
+public class Poisson extends Ones {
+
+  private static final String FUNCTION_NAME = "tf.random.poisson()";
+
+  enum Parameters {
+    SHAPE,
+    LAM,
+    DTYPE,
+    SEED,
+    NAME
+  }
+
+  public Poisson(PointsToSetVariable source, CGNode node) {
+    super(source, node);
+  }
+
+  @Override
+  protected int getDTypeParameterPosition() {
+    return DTYPE.ordinal();
+  }
+
+  protected int getLamParameterPosition() {
+    return LAM.ordinal();
+  }
+
+  protected int getLamParameterValueNumber(PropagationCallGraphBuilder builder) {
+    Set<Integer> numberOfPossiblePositionalArguments =
+        this.getNumberOfPossiblePositionalArguments(builder);
+    int lamParameterPosition = this.getLamParameterPosition();
+
+    if (!numberOfPossiblePositionalArguments.stream().anyMatch(n -> n >= lamParameterPosition + 1))
+      throw new IllegalStateException(
+          "Cannot determine value number for 'lam' parameter of " + FUNCTION_NAME);
+
+    return this.getArgumentValueNumber(lamParameterPosition);
+  }
+
+  @Override
+  protected Set<List<Dimension<?>>> getShapes(PropagationCallGraphBuilder builder) {
+    Set<List<Dimension<?>>> ret = HashSetFactory.make();
+    Set<List<Dimension<?>>> shapes = super.getShapes(builder);
+
+    if (shapes.isEmpty())
+      throw new IllegalStateException(
+          "Cannot determine shape for " + this.getSignature() + " call.");
+
+    // Get the shape of the alpha parameter.
+    Set<List<Dimension<?>>> lamShapes =
+        this.getShapes(builder, this.getLamParameterValueNumber(builder));
+
+    // return shape `tf.concat([shape, tf.shape(lam)], axis=0)`.
+    shapes.forEach(
+        shape -> {
+          lamShapes.forEach(
+              lShape -> {
+                List<Dimension<?>> newShape = new ArrayList<>(shape);
+                newShape.addAll(lShape);
+                ret.add(newShape);
+              });
+        });
+
+    return ret;
+  }
+
+  @Override
+  protected String getSignature() {
+    return FUNCTION_NAME;
+  }
+}

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

@@ -8,6 +8,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.ONE_HOT;
+import static com.ibm.wala.cast.python.ml.types.TensorFlowTypes.POISSON;
 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;
@@ -56,6 +57,7 @@ else if (calledFunction.equals(CONVERT_TO_TENSOR.getDeclaringClass()))
     else if (calledFunction.equals(ONE_HOT.getDeclaringClass())) return new OneHot(source, node);
     else if (calledFunction.equals(EYE.getDeclaringClass())) return new Eye(source, node);
     else if (calledFunction.equals(GAMMA.getDeclaringClass())) return new Gamma(source, node);
+    else if (calledFunction.equals(POISSON.getDeclaringClass())) return new Poisson(source, node);
     else
       throw new IllegalArgumentException(
           "Unknown call: " + calledFunction + " for source: " + source + ".");

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

@@ -164,6 +164,13 @@ public boolean canConvertTo(DType other) {
               PythonTypes.pythonLoader, TypeName.string2TypeName("Ltensorflow/functions/gamma")),
           AstMethodReference.fnSelector);
 
+  /** https://www.tensorflow.org/api_docs/python/tf/poisson. */
+  public static final MethodReference POISSON =
+      MethodReference.findOrCreate(
+          TypeReference.findOrCreate(
+              PythonTypes.pythonLoader, TypeName.string2TypeName("Ltensorflow/functions/poisson")),
+          AstMethodReference.fnSelector);
+
   /**
    * Represents the TensorFlow float32 data type.
    *

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

@@ -0,0 +1,14 @@
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+samples = tf.random.poisson([10], [0.5, 1.5])
+# samples has shape [10, 2], where each slice [:, 0] and [:, 1] represents
+# the samples drawn from each distribution
+assert samples.shape == (10, 2)
+assert samples.dtype == tf.float32
+
+f(samples)

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

@@ -0,0 +1,14 @@
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+samples = tf.random.poisson([10], [0.5, 1.5], tf.float32)
+# samples has shape [10, 2], where each slice [:, 0] and [:, 1] represents
+# the samples drawn from each distribution
+assert samples.shape == (10, 2)
+assert samples.dtype == tf.float32
+
+f(samples)

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

@@ -0,0 +1,14 @@
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+samples = tf.random.poisson([10], [0.5, 1.5], tf.float64)
+# samples has shape [10, 2], where each slice [:, 0] and [:, 1] represents
+# the samples drawn from each distribution
+assert samples.shape == (10, 2)
+assert samples.dtype == tf.float32
+
+f(samples)

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

@@ -0,0 +1,14 @@
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+samples = tf.random.poisson([7, 5], [12.2, 3.3])
+# samples has shape [7, 5, 2], where each slice [:, :, 0] and [:, :, 1]
+# represents the 7x5 samples drawn from each of the two distributions
+assert samples.shape == (7, 5, 2)
+assert samples.dtype == tf.float32
+
+f(samples)