Progress.

Author: khatchad

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

@@ -95,6 +95,12 @@ public class TestTensorflow2Model extends TestPythonMLCallGraphShape {
   private static final TensorType TENSOR_3_3_INT32 =
       new TensorType(INT_32, asList(new NumericDim(3), new NumericDim(3)));
 
+  private static final TensorType TENSOR_0_NONE_FLOAT32 =
+      new TensorType(FLOAT_32, asList(new NumericDim(0), null));
+
+  private static final TensorType TENSOR_0_NONE_3_FLOAT32 =
+      new TensorType(FLOAT_32, asList(new NumericDim(0), null, new NumericDim(3)));
+
   @SuppressWarnings("unused")
   private static final TensorType TENSOR_1_NONE_INT32 =
       new TensorType(INT_32, asList(new NumericDim(1), null));
@@ -111,6 +117,9 @@ public class TestTensorflow2Model extends TestPythonMLCallGraphShape {
   private static final TensorType TENSOR_2_NONE_2_FLOAT32 =
       new TensorType(FLOAT_32, asList(new NumericDim(2), null, new NumericDim(2)));
 
+  private static final TensorType TENSOR_2_NONE_2_INT32 =
+      new TensorType(INT_32, asList(new NumericDim(2), null, new NumericDim(2)));
+
   @SuppressWarnings("unused")
   private static final TensorType TENSOR_2_NONE_NONE_NONE_INT32 =
       new TensorType(INT_32, asList(new NumericDim(2), null));
@@ -4668,6 +4677,21 @@ public void testRaggedConstant12() throws ClassHierarchyException, CancelExcepti
     test("tf2_test_ragged_constant12.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_NONE_2_FLOAT32)));
   }
 
+  @Test
+  public void testRaggedConstant13() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_ragged_constant13.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_0_NONE_FLOAT32)));
+  }
+
+  @Test
+  public void testRaggedConstant14() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_ragged_constant14.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_0_NONE_3_FLOAT32)));
+  }
+
+  @Test
+  public void testRaggedConstant15() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_ragged_constant15.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_NONE_2_INT32)));
+  }
+
   private void test(
       String filename,
       String functionName,

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

@@ -1,12 +1,15 @@
 package com.ibm.wala.cast.python.ml.client;
 
+import static com.ibm.wala.cast.python.ml.client.RaggedConstant.Parameters.INNER_SHAPE;
 import static com.ibm.wala.cast.python.ml.client.RaggedConstant.Parameters.RAGGED_RANK;
 import static com.ibm.wala.cast.python.ml.types.TensorFlowTypes.DType.FLOAT32;
 import static com.ibm.wala.cast.python.types.PythonTypes.Root;
 import static com.ibm.wala.cast.python.types.PythonTypes.list;
 import static com.ibm.wala.cast.python.types.PythonTypes.tuple;
 import static com.ibm.wala.cast.python.util.Util.getAllocationSiteInNode;
 import static com.ibm.wala.core.util.strings.Atom.findOrCreateAsciiAtom;
+import static java.lang.Math.max;
+import static java.util.Collections.emptySet;
 import static java.util.logging.Logger.getLogger;
 import static java.util.stream.Collectors.toSet;
 
@@ -306,8 +309,14 @@ protected Set<List<Dimension<?>>> getShapesOfValue(
       LOGGER.fine("Tensor rank: " + K);
 
       Set<Long> rankArguments = this.getPossibleRaggedRankArguments(builder);
+      Set<List<Dimension<?>>> innerShapeArguments = this.getPossibleInnerShapeArguments(builder);
 
-      if (rankArguments.isEmpty()) rankArguments.add(K - 1L); // Default ragged rank.
+      if (rankArguments.isEmpty())
+        // Default ragged rank.
+        if (innerShapeArguments.isEmpty()) rankArguments.add(max(0, K - 1L));
+        else
+          for (List<Dimension<?>> innerShape : innerShapeArguments)
+            rankArguments.add(max(0, K - 1L - innerShape.size()));
 
       for (Long R : rankArguments) {
         LOGGER.fine("Ragged rank: " + R);
@@ -375,36 +384,7 @@ private static int getMaximumDepthOfInstance(
   }
 
   protected Set<Long> getPossibleRaggedRankArguments(PropagationCallGraphBuilder builder) {
-    Set<Long> ret = HashSetFactory.make();
-    int valueNumber = this.getRaggedRankArgumentValueNumber(builder);
-
-    if (valueNumber >= 0) {
-      PointerAnalysis<InstanceKey> pointerAnalysis = builder.getPointerAnalysis();
-      PointerKey raggedRankPK =
-          pointerAnalysis.getHeapModel().getPointerKeyForLocal(this.getNode(), valueNumber);
-      OrdinalSet<InstanceKey> raggedRankPointsToSet = pointerAnalysis.getPointsToSet(raggedRankPK);
-
-      if (raggedRankPointsToSet == null || raggedRankPointsToSet.isEmpty())
-        throw new IllegalArgumentException(
-            "Empty points-to set for ragged_rank in source: " + this.getSource() + ".");
-
-      for (InstanceKey raggedRankIK : raggedRankPointsToSet)
-        if (raggedRankIK instanceof ConstantKey) {
-          ConstantKey<?> constantKey = (ConstantKey<?>) raggedRankIK;
-          Object constantKeyValue = constantKey.getValue();
-
-          if (constantKeyValue instanceof Long) {
-            Long raggedRankValue = (Long) constantKeyValue;
-            ret.add(raggedRankValue);
-          } else
-            throw new IllegalArgumentException(
-                "Expected an integer for ragged_rank, but found: " + constantKeyValue + ".");
-        } else
-          throw new IllegalArgumentException(
-              "Expected a constant key for ragged_rank, but found: " + raggedRankIK + ".");
-    }
-
-    return ret;
+    return this.getPossibleLongArguments(builder, this.getRaggedRankArgumentValueNumber(builder));
   }
 
   protected int getRaggedRankParameterPosition() {
@@ -416,6 +396,26 @@ protected int getRaggedRankArgumentValueNumber(PropagationCallGraphBuilder build
     return this.getArgumentValueNumber(builder, this.getRaggedRankParameterPosition(), true);
   }
 
+  protected int getInnerShapeParameterPosition() {
+    return INNER_SHAPE.ordinal();
+  }
+
+  protected int getInnerShapeArgumentValueNumber(PropagationCallGraphBuilder builder) {
+    // TODO: Handle keyword arguments.
+    return this.getArgumentValueNumber(builder, this.getInnerShapeParameterPosition(), true);
+  }
+
+  protected Set<List<Dimension<?>>> getPossibleInnerShapeArguments(
+      PropagationCallGraphBuilder builder) {
+    int valueNumber = this.getInnerShapeArgumentValueNumber(builder);
+
+    if (valueNumber >= 0) {
+      PointerKey pointerKey = builder.getPointerKeyForLocal(this.getNode(), valueNumber);
+      OrdinalSet<InstanceKey> pointsToSet = builder.getPointerAnalysis().getPointsToSet(pointerKey);
+      return this.getShapesFromShapeArgument(builder, pointsToSet);
+    } else return emptySet();
+  }
+
   /**
    * {@inheritDoc}
    *

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

@@ -88,7 +88,7 @@ public Set<TensorType> getTensorTypes(PropagationCallGraphBuilder builder) {
    * Returns the possible shapes of the tensor returned by this generator.
    *
    * @param builder The {@link PropagationCallGraphBuilder} used to build the call graph.
-   * @param pointsToSet The points-to set of the shape argument.
+   * @param pointsToSet The points-to set of the shape argument. FIXME: Why not take a value number?
    * @return A set of possible shapes of the tensor returned by this generator.
    */
   protected Set<List<Dimension<?>>> getShapesFromShapeArgument(
@@ -290,6 +290,7 @@ protected Set<List<Dimension<?>>> getShapes(
       throw new IllegalArgumentException(
           "Empty points-to set for value number: " + valueNumber + " in: " + this.getNode() + ".");
 
+    // FIXME: Just use the value number directly?
     return getShapesOfValue(builder, valuePointsToSet);
   }
 
@@ -736,4 +737,37 @@ protected Set<Integer> getNumberOfPossiblePositionalArguments(
 
     return ret;
   }
+
+  protected Set<Long> getPossibleLongArguments(
+      PropagationCallGraphBuilder builder, int valueNumber) {
+    Set<Long> ret = HashSetFactory.make();
+
+    if (valueNumber >= 0) {
+      PointerAnalysis<InstanceKey> pointerAnalysis = builder.getPointerAnalysis();
+      PointerKey pointerKey =
+          pointerAnalysis.getHeapModel().getPointerKeyForLocal(this.getNode(), valueNumber);
+      OrdinalSet<InstanceKey> pointsToSet = pointerAnalysis.getPointsToSet(pointerKey);
+
+      if (pointsToSet == null || pointsToSet.isEmpty())
+        throw new IllegalArgumentException(
+            "Empty points-to set in source: " + this.getSource() + ".");
+
+      for (InstanceKey instanceKey : pointsToSet)
+        if (instanceKey instanceof com.ibm.wala.ipa.callgraph.propagation.ConstantKey) {
+          ConstantKey<?> constantKey = (ConstantKey<?>) instanceKey;
+          Object constantKeyValue = constantKey.getValue();
+
+          if (constantKeyValue instanceof Long) {
+            Long value = (Long) constantKeyValue;
+            ret.add(value);
+          } else
+            throw new IllegalStateException(
+                "Expected a long, but found: " + constantKeyValue + ".");
+        } else
+          throw new IllegalStateException(
+              "Expected a constant key, but found: " + instanceKey + ".");
+    }
+
+    return ret;
+  }
 }

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

@@ -1,11 +1,6 @@
 package com.ibm.wala.cast.python.ml.client;
 
-import com.ibm.wala.cast.python.ml.types.TensorType.Dimension;
-import com.ibm.wala.ipa.callgraph.propagation.InstanceKey;
 import com.ibm.wala.ipa.callgraph.propagation.PointsToSetVariable;
-import com.ibm.wala.ipa.callgraph.propagation.PropagationCallGraphBuilder;
-import java.util.List;
-import java.util.Set;
 
 /**
  * A generator for tensors created by the `zeros_like()` function in TensorFlow.
@@ -28,13 +23,6 @@ public ZerosLike(PointsToSetVariable source) {
     super(source);
   }
 
-  @Override
-  protected Set<List<Dimension<?>>> getShapesFromShapeArgument(
-      PropagationCallGraphBuilder builder, Iterable<InstanceKey> pointsToSet) {
-    throw new UnsupportedOperationException(
-        "Shapes are derived from the `input` argument and cannot be provided explicitly.");
-  }
-
   @Override
   protected int getShapeParameterPosition() {
     return SHAPE_PARAMETER_POSITION;

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

@@ -0,0 +1,16 @@
+# From https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/GradientTape#gradient.
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+# TensorFlow sees an empty list and gives up on the inner dimensions.
+t = tf.ragged.constant([], None, 1)
+assert t.shape == (0, None)
+assert t.dtype == tf.float32
+# Output: (0, None) -> It lost the inner dimension info!
+
+f(t)

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

@@ -0,0 +1,17 @@
+# From https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/GradientTape#gradient.
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+# You tell TensorFlow: "Even though it's empty, if there WERE data,
+# it would be shape (3,)."
+t = tf.ragged.constant([], None, 1, (3,))
+assert t.shape == (0, None, 3)
+assert t.dtype == tf.float32
+# Output: (0, None, 3) -> The inner structure is preserved.
+
+f(t)

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

@@ -0,0 +1,22 @@
+# From https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/GradientTape#gradient.
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+data = [[[1, 2], [3, 4]], [[5, 6]]]
+
+# Success: The data matches the inner shape (2,)
+t1 = tf.ragged.constant(data, None, 1, (2,))
+assert t1.shape == (2, None, 2)
+assert t1.dtype == tf.int32
+# Output: (2, None, 2)
+
+# Failure: You claim inner shape is (3,), but data is length 2
+# t2 = tf.ragged.constant(data, ragged_rank=1, inner_shape=(3,))
+# Raises ValueError: Inner shape mismatch
+
+f(t1)