Merge branch '267-initial-tensor-dimensions-arent-always-accurate' of https://github.com/ponder-lab/ML into 267-initial-tensor-dimensions-arent-always-accurate

Author: khatchad

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

@@ -21,6 +21,7 @@
 import com.ibm.wala.cast.python.ipa.callgraph.PythonSSAPropagationCallGraphBuilder;
 import com.ibm.wala.cast.python.ml.analysis.TensorTypeAnalysis;
 import com.ibm.wala.cast.python.ml.analysis.TensorVariable;
+import com.ibm.wala.cast.python.ml.client.NonBroadcastableShapesException;
 import com.ibm.wala.cast.python.ml.client.PythonTensorAnalysisEngine;
 import com.ibm.wala.cast.python.ml.types.TensorType;
 import com.ibm.wala.cast.python.ml.types.TensorType.Dimension;
@@ -120,6 +121,12 @@ public class TestTensorflow2Model extends TestPythonMLCallGraphShape {
   private static final TensorType TENSOR_2_NONE_2_INT32 =
       new TensorType(INT_32, asList(new NumericDim(2), null, new NumericDim(2)));
 
+  private static final TensorType TENSOR_2_NONE_2_3_INT32 =
+      new TensorType(INT_32, asList(new NumericDim(2), null, new NumericDim(2), new NumericDim(3)));
+
+  private static final TensorType TENSOR_2_NONE_2_2_INT32 =
+      new TensorType(INT_32, asList(new NumericDim(2), null, new NumericDim(2), new NumericDim(2)));
+
   @SuppressWarnings("unused")
   private static final TensorType TENSOR_2_NONE_NONE_NONE_INT32 =
       new TensorType(INT_32, asList(new NumericDim(2), null));
@@ -186,6 +193,11 @@ public class TestTensorflow2Model extends TestPythonMLCallGraphShape {
           FLOAT_32,
           asList(new NumericDim(3), new NumericDim(2), new NumericDim(2), new NumericDim(3)));
 
+  private static final TensorType TENSOR_2_2_2_3_FLOAT32 =
+      new TensorType(
+          FLOAT_32,
+          asList(new NumericDim(2), new NumericDim(2), new NumericDim(2), new NumericDim(3)));
+
   private static final TensorType TENSOR_20_28_28_FLOAT32 =
       new TensorType(FLOAT_32, asList(new NumericDim(20), new NumericDim(28), new NumericDim(28)));
 
@@ -227,7 +239,11 @@ public void testValueIndex()
         Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
   }
 
-  @Test
+  /**
+   * Should not throw an {@link IllegalArgumentException} once https://github.com/wala/ML/issues/340
+   * is fixed.
+   */
+  @Test(expected = IllegalArgumentException.class)
   public void testValueIndex2()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
     test(
@@ -238,7 +254,11 @@ public void testValueIndex2()
         Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
   }
 
-  @Test
+  /**
+   * Should not throw an {@link IllegalArgumentException} once https://github.com/wala/ML/issues/340
+   * is fixed.
+   */
+  @Test(expected = IllegalArgumentException.class)
   public void testValueIndex3()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
     test(
@@ -249,7 +269,11 @@ public void testValueIndex3()
         Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
   }
 
-  @Test
+  /**
+   * Should not throw an {@link IllegalArgumentException} once https://github.com/wala/ML/issues/340
+   * is fixed.
+   */
+  @Test(expected = IllegalArgumentException.class)
   public void testValueIndex4()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
     test(
@@ -260,7 +284,11 @@ public void testValueIndex4()
         Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
   }
 
-  @Test
+  /**
+   * Should not throw an {@link IllegalArgumentException} once https://github.com/wala/ML/issues/340
+   * is fixed.
+   */
+  @Test(expected = IllegalArgumentException.class)
   public void testFunction()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
     test("tf2_test_function.py", "func2", 1, 1, Map.of(2, Set.of(MNIST_INPUT)));
@@ -1617,25 +1645,45 @@ public void testAdd58()
   @Test
   public void testAdd59()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
-    test("tf2_test_add59.py", "add", 2, 2, Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
+    test(
+        "tf2_test_add59.py",
+        "add",
+        2,
+        2,
+        Map.of(2, Set.of(TENSOR_2_INT32), 3, Set.of(TENSOR_2_INT32)));
   }
 
   @Test
   public void testAdd60()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
-    test("tf2_test_add60.py", "add", 2, 2, Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
+    test(
+        "tf2_test_add60.py",
+        "add",
+        2,
+        2,
+        Map.of(2, Set.of(TENSOR_2_INT32), 3, Set.of(TENSOR_2_INT32)));
   }
 
   @Test
   public void testAdd61()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
-    test("tf2_test_add61.py", "add", 2, 2, Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
+    test(
+        "tf2_test_add61.py",
+        "add",
+        2,
+        2,
+        Map.of(2, Set.of(TENSOR_2_INT32), 3, Set.of(TENSOR_2_INT32)));
   }
 
   @Test
   public void testAdd62()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
-    test("tf2_test_add62.py", "add", 2, 2, Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
+    test(
+        "tf2_test_add62.py",
+        "add",
+        2,
+        2,
+        Map.of(2, Set.of(TENSOR_2_INT32), 3, Set.of(TENSOR_2_INT32)));
   }
 
   @Test
@@ -2151,16 +2199,24 @@ public void testReduceMean3()
   @Test
   public void testGradient()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
-    test("tf2_test_gradient.py", "f", 1, 1, Map.of(2, Set.of(MNIST_INPUT)));
+    test("tf2_test_gradient.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_NONE_FLOAT32)));
   }
 
-  @Test
+  /**
+   * Should not throw an {@link IllegalArgumentException} once https://github.com/wala/ML/issues/340
+   * is fixed.
+   */
+  @Test(expected = IllegalArgumentException.class)
   public void testGradient2()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
     test("tf2_test_gradient2.py", "f", 1, 1, Map.of(2, Set.of(MNIST_INPUT)));
   }
 
-  @Test
+  /**
+   * Should not throw an {@link IllegalArgumentException} once https://github.com/wala/ML/issues/340
+   * is fixed.
+   */
+  @Test(expected = IllegalArgumentException.class)
   public void testMultiply()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
     test("tf2_test_multiply.py", "f", 1, 1, Map.of(2, Set.of(MNIST_INPUT)));
@@ -2169,7 +2225,48 @@ public void testMultiply()
   @Test
   public void testMultiply2()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
-    test("tf2_test_multiply2.py", "f", 1, 1, Map.of(2, Set.of(MNIST_INPUT)));
+    test("tf2_test_multiply2.py", "f", 1, 1, Map.of(2, Set.of(SCALAR_TENSOR_OF_INT32)));
+  }
+
+  @Test
+  public void testMultiply3()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_multiply3.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_3_FLOAT32)));
+  }
+
+  @Test
+  public void testMultiply4()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_multiply4.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_3_FLOAT32)));
+  }
+
+  @Test
+  public void testMultiply5()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_multiply5.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_2_2_3_FLOAT32)));
+  }
+
+  /**
+   * This is an invalid case since the inputs have different ranks.
+   *
+   * <p>For now, we are throwing an exception. But, this is invalid code.
+   *
+   * <p>TODO: We'll need to come up with a suitable way to handle this in the future.
+   */
+  @Test(expected = NonBroadcastableShapesException.class)
+  public void testMultiply6()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_multiply6.py", "f", 1, 1);
+  }
+
+  /**
+   * Should not throw an {@link IllegalArgumentException} once https://github.com/wala/ML/issues/340
+   * is fixed.
+   */
+  @Test(expected = IllegalArgumentException.class)
+  public void testMultiply7()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_multiply7.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_3_FLOAT32)));
   }
 
   @Test
@@ -4697,6 +4794,22 @@ public void testRaggedConstant16() throws ClassHierarchyException, CancelExcepti
     test("tf2_test_ragged_constant16.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_NONE_2_INT32)));
   }
 
+  /**
+   * Test non-uniform inner dimensions.
+   *
+   * <p>TODO: Remove expected assertion error once https://github.com/wala/ML/issues/350 is fixed.
+   */
+  @Test(expected = AssertionError.class)
+  public void testRaggedConstant17() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_ragged_constant17.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_NONE_2_3_INT32)));
+  }
+
+  /** This one works because the inner dimensions are uniform. */
+  @Test
+  public void testRaggedConstant18() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_ragged_constant18.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_NONE_2_2_INT32)));
+  }
+
   private void test(
       String filename,
       String functionName,

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

@@ -32,7 +32,7 @@ public Constant(PointsToSetVariable source) {
   protected Set<List<Dimension<?>>> getDefaultShapes(PropagationCallGraphBuilder builder) {
     // If the shape argument is not specified, then the shape is inferred from the shape of value.
     // TODO: Handle keyword arguments.
-    return getShapes(builder, this.getValueArgumentValueNumber());
+    return this.getShapes(builder, this.getValueArgumentValueNumber());
   }
 
   /**

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

@@ -0,0 +1,102 @@
+package com.ibm.wala.cast.python.ml.client;
+
+import static com.ibm.wala.cast.python.ml.client.Multiply.Parameters.X;
+import static com.ibm.wala.cast.python.ml.client.Multiply.Parameters.Y;
+import static com.ibm.wala.cast.python.ml.types.TensorFlowTypes.TYPE_REFERENCE_TO_SIGNATURE;
+import static com.ibm.wala.cast.python.ml.util.TensorShapeUtil.areBroadcastable;
+import static com.ibm.wala.cast.python.ml.util.TensorShapeUtil.getBroadcastedShapes;
+import static com.ibm.wala.cast.python.util.Util.getFunction;
+import static java.util.logging.Logger.getLogger;
+
+import com.ibm.wala.cast.python.ml.types.TensorType.Dimension;
+import com.ibm.wala.ipa.callgraph.propagation.PointsToSetVariable;
+import com.ibm.wala.ipa.callgraph.propagation.PropagationCallGraphBuilder;
+import com.ibm.wala.types.TypeReference;
+import com.ibm.wala.util.collections.HashSetFactory;
+import java.util.List;
+import java.util.Set;
+import java.util.logging.Logger;
+
+/**
+ * A representation of a multiply operation in TensorFlow.
+ *
+ * @see <a href="https://www.tensorflow.org/api_docs/python/tf/multiply">tf.multiply</a>.
+ * @author <a href="mailto:[email protected]">Raffi Khatchadourian</a>
+ */
+public class Multiply extends ZerosLike {
+
+  @SuppressWarnings("unused")
+  private static final Logger logger = getLogger(Multiply.class.getName());
+
+  protected enum Parameters {
+    X,
+    Y,
+    NAME
+  }
+
+  /**
+   * The dtype argument is not explicitly provided to multiply(); rather, the dtype is inferred from
+   * the `x` argument.
+   *
+   * @see <a
+   *     href="https://www.tensorflow.org/api_docs/python/tf/math/multiply#returns">tf.math.multiply
+   *     - Returns</a>.
+   */
+  protected static final int DTYPE_PARAMETER_POSITION = -1;
+
+  @Override
+  protected int getDTypeParameterPosition() {
+    return DTYPE_PARAMETER_POSITION;
+  }
+
+  public Multiply(PointsToSetVariable source) {
+    super(source);
+  }
+
+  protected int getXParameterPosition() {
+    return X.ordinal();
+  }
+
+  protected int getXArgumentValueNumber(PropagationCallGraphBuilder builder) {
+    // TODO: Handle keyword arguments.
+    return this.getArgumentValueNumber(builder, this.getXParameterPosition());
+  }
+
+  protected int getYParameterPosition() {
+    return Y.ordinal();
+  }
+
+  protected int getYArgumentValueNumber(PropagationCallGraphBuilder builder) {
+    // TODO: Handle keyword arguments.
+    return this.getArgumentValueNumber(builder, this.getYParameterPosition());
+  }
+
+  /**
+   * Returns the TensorFlow function signature represented by this generator.
+   *
+   * @return The TensorFlow function signature represented by this generator.
+   */
+  @Override
+  protected String getSignature() {
+    TypeReference function = getFunction(this.getSource());
+    return TYPE_REFERENCE_TO_SIGNATURE.get(function);
+  }
+
+  @Override
+  protected Set<List<Dimension<?>>> getDefaultShapes(PropagationCallGraphBuilder builder) {
+    // The resulting shape is the broadcasted shape of the shapes of x and y.
+    Set<List<Dimension<?>>> ret = HashSetFactory.make();
+
+    Set<List<Dimension<?>>> xShapes =
+        this.getShapes(builder, this.getXArgumentValueNumber(builder));
+    Set<List<Dimension<?>>> yShapes =
+        this.getShapes(builder, this.getYArgumentValueNumber(builder));
+
+    for (List<Dimension<?>> xShape : xShapes)
+      for (List<Dimension<?>> yShape : yShapes)
+        if (areBroadcastable(xShape, yShape)) ret.add(getBroadcastedShapes(xShape, yShape));
+        else throw new NonBroadcastableShapesException(this, xShape, yShape);
+
+    return ret;
+  }
+}

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

@@ -0,0 +1,47 @@
+package com.ibm.wala.cast.python.ml.client;
+
+import static java.lang.String.format;
+
+import com.ibm.wala.cast.python.ml.types.TensorType.Dimension;
+import java.util.List;
+
+/**
+ * An exception indicating that two shapes are not broadcastable for a given operation.
+ *
+ * @author <a href="mailto:[email protected]">Raffi Khatchadourian</a>
+ * @see <a href="https://numpy.org/doc/stable/user/basics.broadcasting.html">NumPy Broadcasting</a>.
+ */
+public class NonBroadcastableShapesException extends RuntimeException {
+
+  /** Serial version UID. */
+  private static final long serialVersionUID = 805036824027449575L;
+
+  /** The operation for which the shapes are not broadcastable. */
+  private final transient Object op;
+
+  /** The first shape. */
+  private final transient List<Dimension<?>> xShape;
+
+  /** The second shape. */
+  private final transient List<Dimension<?>> yShape;
+
+  /**
+   * Constructs a new exception indicating that the given shapes are not broadcastable for the given
+   * operation.
+   *
+   * @param op The operation for which the shapes are not broadcastable.
+   * @param xShape The first shape.
+   * @param yShape The second shape.
+   */
+  public NonBroadcastableShapesException(
+      Object op, List<Dimension<?>> xShape, List<Dimension<?>> yShape) {
+    this.op = op;
+    this.xShape = xShape;
+    this.yShape = yShape;
+  }
+
+  @Override
+  public String getMessage() {
+    return format("The shapes %s and %s are not broadcastable for %s.", xShape, yShape, op);
+  }
+}