Handle batch dimensions for tf.eye().

Author: khatchad

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

@@ -110,6 +110,14 @@ public class TestTensorflow2Model extends TestPythonMLCallGraphShape {
   private static final TensorType TENSOR_2_5_3_FLOAT32 =
       new TensorType(FLOAT_32, asList(new NumericDim(2), new NumericDim(5), new NumericDim(3)));
 
+  private static final TensorType TENSOR_3_2_2_FLOAT32 =
+      new TensorType(FLOAT_32, asList(new NumericDim(3), new NumericDim(2), new NumericDim(2)));
+
+  private static final TensorType TENSOR_3_2_2_3_FLOAT32 =
+      new TensorType(
+          FLOAT_32,
+          asList(new NumericDim(3), 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)));
 
@@ -4370,6 +4378,18 @@ public void testEye3()
     test("tf2_test_eye3.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_3_FLOAT32)));
   }
 
+  @Test
+  public void testEye4()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_eye4.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_3_2_2_FLOAT32)));
+  }
+
+  @Test
+  public void testEye5()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_eye5.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_3_2_2_3_FLOAT32)));
+  }
+
   private void test(
       String filename,
       String functionName,

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

@@ -4,6 +4,7 @@
 import static com.ibm.wala.cast.python.ml.client.Eye.Parameters.DTYPE;
 import static com.ibm.wala.cast.python.ml.client.Eye.Parameters.NUM_COLUMNS;
 import static com.ibm.wala.cast.python.ml.client.Eye.Parameters.NUM_ROWS;
+import static java.util.Collections.emptySet;
 
 import com.ibm.wala.cast.python.ml.types.TensorType.Dimension;
 import com.ibm.wala.cast.python.ml.types.TensorType.NumericDim;
@@ -14,6 +15,7 @@
 import com.ibm.wala.ipa.callgraph.propagation.PointsToSetVariable;
 import com.ibm.wala.ipa.callgraph.propagation.PropagationCallGraphBuilder;
 import com.ibm.wala.util.collections.HashSetFactory;
+import com.ibm.wala.util.intset.OrdinalSet;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
@@ -53,26 +55,27 @@ protected int getNumRowsParameterPosition() {
     return NUM_ROWS.ordinal();
   }
 
-  protected int getNumRowsValueNumber(PropagationCallGraphBuilder builder) {
+  protected int getNumRowsArgumentValueNumber() {
     return this.getArgumentValueNumber(this.getNumRowsParameterPosition());
   }
 
+  protected int getBatchShapesArgumentValueNumber() {
+    // TOOD: Handle keyword arguments.
+    return this.getArgumentValueNumber(this.getBatchShapeParameterPosition());
+  }
+
   protected int getNumColumnsParameterPosition() {
     return NUM_COLUMNS.ordinal();
   }
 
-  protected int getNumColumnsValueNumber(PropagationCallGraphBuilder builder) {
+  protected int getNumColumnsArgumentValueNumber(PropagationCallGraphBuilder builder) {
     return this.getArgumentValueNumber(this.getNumColumnsParameterPosition());
   }
 
   protected int getBatchShapeParameterPosition() {
     return BATCH_SHAPE.ordinal();
   }
 
-  protected int getBatchShapeValueNumber(PropagationCallGraphBuilder builder) {
-    return this.getArgumentValueNumber(this.getBatchShapeParameterPosition());
-  }
-
   @Override
   protected Set<List<Dimension<?>>> getShapes(PropagationCallGraphBuilder builder) {
     Set<List<Dimension<?>>> ret = HashSetFactory.make();
@@ -114,6 +117,12 @@ protected Set<List<Dimension<?>>> getShapes(PropagationCallGraphBuilder builder)
         }
     }
 
+    Set<List<Dimension<?>>> batchShapes = this.getBatchShapes(builder);
+
+    // prepend batch dimensions to each shape.
+    for (List<Dimension<?>> batchDim : batchShapes)
+      for (List<Dimension<?>> retDim : ret) retDim.addAll(0, batchDim);
+
     return ret;
   }
 
@@ -133,6 +142,33 @@ private Set<Optional<Integer>> getNumberOfColumns(PropagationCallGraphBuilder bu
     return this.getPossiblePositionalArgumentValues(builder, this.getNumColumnsParameterPosition());
   }
 
+  private Set<List<Dimension<?>>> getBatchShapes(PropagationCallGraphBuilder builder) {
+    // TODO Handle keyword arguments.
+    Set<Integer> possibleNumArgs = this.getNumberOfPossiblePositionalArguments(builder);
+
+    if (possibleNumArgs.contains(this.getBatchShapeParameterPosition() + 1)) {
+      PointerAnalysis<InstanceKey> pointerAnalysis = builder.getPointerAnalysis();
+
+      PointerKey pointerKey =
+          pointerAnalysis
+              .getHeapModel()
+              .getPointerKeyForLocal(this.getNode(), this.getBatchShapesArgumentValueNumber());
+
+      OrdinalSet<InstanceKey> pts = pointerAnalysis.getPointsToSet(pointerKey);
+
+      Set<List<Dimension<?>>> shapesFromShapeArgument =
+          this.getShapesFromShapeArgument(builder, pts);
+
+      if (shapesFromShapeArgument == null || shapesFromShapeArgument.isEmpty())
+        throw new IllegalStateException(
+            "Batch shape argument for tf.eye() should be a list of dimensions.");
+
+      return shapesFromShapeArgument;
+    }
+
+    return emptySet();
+  }
+
   private Set<Optional<Integer>> getPossiblePositionalArgumentValues(
       PropagationCallGraphBuilder builder, int paramPosition) {
     PointerAnalysis<InstanceKey> pointerAnalysis = builder.getPointerAnalysis();

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

@@ -103,7 +103,7 @@ protected Set<List<Dimension<?>>> getShapesFromShapeArgument(
       AllocationSiteInNode asin = getAllocationSiteInNode(instanceKey);
       TypeReference reference = asin.getConcreteType().getReference();
 
-      if (reference.equals(list)) { // TODO: This can also be a tuple of tensors.
+      if (reference.equals(list) || reference.equals(tuple)) {
         // We have a list of integers that represent the shape.
         OrdinalSet<InstanceKey> objectCatalogPointsToSet =
             pointerAnalysis.getPointsToSet(
@@ -206,7 +206,13 @@ protected Set<List<Dimension<?>>> getShapesFromShapeArgument(
           }
       } else
         throw new IllegalStateException(
-            "Expected a " + PythonTypes.list + " for the shape, but got: " + reference + ".");
+            "Expected a "
+                + PythonTypes.list
+                + " or "
+                + PythonTypes.tuple
+                + " for the shape, but got: "
+                + reference
+                + ".");
     }
 
     return ret;

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

@@ -0,0 +1,14 @@
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+# Construct a batch of 3 identity matrices, each 2 x 2.
+# batch_identity[i, :, :] is a 2 x 2 identity matrix, i = 0, 1, 2.
+arg = tf.eye(2, None, [3])
+assert arg.shape == (3, 2, 2)
+assert arg.dtype == tf.float32
+
+f(arg)

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

@@ -0,0 +1,13 @@
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+arg = tf.eye(2, 3, [3, 2])
+
+assert arg.shape == (3, 2, 2, 3)
+assert arg.dtype == tf.float32
+
+f(arg)