Handle the multiplication of tensors with broadcasting in TensorFlow.

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;
@@ -192,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)));
 
@@ -2215,7 +2221,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

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);
+  }
+}

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.EYE;
 import static com.ibm.wala.cast.python.ml.types.TensorFlowTypes.FILL;
 import static com.ibm.wala.cast.python.ml.types.TensorFlowTypes.GAMMA;
+import static com.ibm.wala.cast.python.ml.types.TensorFlowTypes.MULTIPLY;
 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;
@@ -54,6 +55,7 @@ else if (calledFunction.equals(CONVERT_TO_TENSOR.getDeclaringClass()))
     else if (calledFunction.equals(POISSON.getDeclaringClass())) return new Poisson(source);
     else if (calledFunction.equals(RAGGED_CONSTANT.getDeclaringClass()))
       return new RaggedConstant(source);
+    else if (calledFunction.equals(MULTIPLY.getDeclaringClass())) return new Multiply(source);
     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

@@ -189,9 +189,21 @@ public boolean canConvertTo(DType other) {
 
   private static final String RAGGED_CONSTANT_SIGNATURE = "tf.ragged.constant()";
 
+  public static final MethodReference MULTIPLY =
+      MethodReference.findOrCreate(
+          TypeReference.findOrCreate(
+              PythonTypes.pythonLoader, TypeName.string2TypeName("Ltensorflow/math/multiply")),
+          AstMethodReference.fnSelector);
+
+  private static final String MULTIPLY_SIGNATURE = "tf.multiply()";
+
   /** A mapping from a {@link TypeReference} to its associated TensorFlow signature. */
   public static final Map<TypeReference, String> TYPE_REFERENCE_TO_SIGNATURE =
-      Map.of(RAGGED_CONSTANT.getDeclaringClass(), RAGGED_CONSTANT_SIGNATURE);
+      Map.of(
+          RAGGED_CONSTANT.getDeclaringClass(),
+          RAGGED_CONSTANT_SIGNATURE,
+          MULTIPLY.getDeclaringClass(),
+          MULTIPLY_SIGNATURE);
 
   /**
    * Represents the TensorFlow float32 data type.

com.ibm.wala.cast.python.ml/source/com/ibm/wala/cast/python/ml/util/TensorShapeUtil.java

@@ -0,0 +1,62 @@
+package com.ibm.wala.cast.python.ml.util;
+
+import static java.lang.Math.max;
+
+import com.ibm.wala.cast.python.ml.types.TensorType.Dimension;
+import com.ibm.wala.cast.python.ml.types.TensorType.NumericDim;
+import java.util.List;
+
+public class TensorShapeUtil {
+
+  public static boolean areBroadcastable(List<Dimension<?>> xShape, List<Dimension<?>> yShape) {
+    int xRank = xShape.size();
+    int yRank = yShape.size();
+    int maxRank = max(xRank, yRank);
+
+    for (int i = 0; i < maxRank; i++) {
+      Dimension<?> xDim = i < (maxRank - xRank) ? null : xShape.get(i - (maxRank - xRank));
+      Dimension<?> yDim = i < (maxRank - yRank) ? null : yShape.get(i - (maxRank - yRank));
+
+      if (xDim == null || yDim == null) {
+        continue; // One of the dimensions is missing, treat as size 1
+      }
+
+      if (xDim instanceof NumericDim && yDim instanceof NumericDim) {
+        int xSize = ((NumericDim) xDim).value();
+        int ySize = ((NumericDim) yDim).value();
+
+        if (xSize != ySize && xSize != 1 && ySize != 1) return false; // Incompatible sizes
+      } else return false; // Non-numeric dimensions are incompatible
+    }
+
+    return true; // All dimensions are compatible
+  }
+
+  public static List<Dimension<?>> getBroadcastedShapes(
+      List<Dimension<?>> xShape, List<Dimension<?>> yShape) {
+    List<Dimension<?>> ret = new java.util.ArrayList<>();
+
+    int xRank = xShape.size();
+    int yRank = yShape.size();
+    int maxRank = max(xRank, yRank);
+
+    for (int i = 0; i < maxRank; i++) {
+      Dimension<?> xDim = i < (maxRank - xRank) ? null : xShape.get(i - (maxRank - xRank));
+      Dimension<?> yDim = i < (maxRank - yRank) ? null : yShape.get(i - (maxRank - yRank));
+
+      if (xDim == null) ret.add(yDim);
+      else if (yDim == null) ret.add(xDim);
+      else if (xDim instanceof NumericDim && yDim instanceof NumericDim) {
+        int xSize = ((NumericDim) xDim).value();
+        int ySize = ((NumericDim) yDim).value();
+
+        if (xSize == ySize) ret.add(xDim); // Both sizes are equal
+        else if (xSize == 1) ret.add(yDim); // x is broadcasted
+        else if (ySize == 1) ret.add(xDim); // y is broadcasted
+        else throw new IllegalArgumentException("Incompatible dimensions for broadcasting.");
+      } else throw new IllegalArgumentException("Non-numeric dimensions cannot be broadcasted.");
+    }
+
+    return ret;
+  }
+}

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

@@ -7,4 +7,8 @@ def f(a):
     pass
 
 
-f(tf.math.multiply(7, 6))
+arg = tf.math.multiply(7, 6)
+assert arg.shape == ()
+assert arg.dtype == tf.int32
+
+f(arg)

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

@@ -0,0 +1,25 @@
+# From https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/math/multiply#for_example/
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+# Shape: (2, 3)
+matrix = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
+assert len(matrix) == 2 and len(matrix[0]) == 3  # Confirming shape (2, 3)
+assert matrix[0][0].__class__ == float  # Confirming dtype float32
+
+# Shape: (1,) -> Broadcasts to (2, 3)
+scalar = [10.0]
+assert len(scalar) == 1  # Confirming shape (1,)
+assert scalar[0].__class__ == float  # Confirming dtype float32
+
+# 1. Scalar Multiplication
+result_scalar = tf.multiply(matrix, scalar)
+assert result_scalar.shape == (2, 3)
+assert result_scalar.dtype == tf.float32
+
+f(result_scalar)

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

@@ -0,0 +1,27 @@
+# From https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/math/multiply#for_example/
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+# Shape: (2, 3)
+matrix = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
+assert len(matrix) == 2 and len(matrix[0]) == 3  # Confirming shape (2, 3)
+assert matrix[0][0].__class__ == float  # Confirming dtype float32
+
+# Shape: (2, 1) -> Broadcasts columns to match matrix width (3)
+col_vector = [[2.0], [3.0]]
+assert len(col_vector) == 2 and len(col_vector[0]) == 1  # Confirming shape (2, 1)
+assert col_vector[0][0].__class__ == float  # Confirming dtype float32
+
+# 2. Column Vector Multiplication
+# [1, 2, 3] * 2 = [2, 4, 6]
+# [4, 5, 6] * 3 = [12, 15, 18]
+result_col = tf.multiply(matrix, col_vector)
+assert result_col.shape == (2, 3)
+assert result_col.dtype == tf.float32
+
+f(result_col)