Progress.

Author: khatchad

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

@@ -105,6 +105,12 @@ public class TestTensorflow2Model extends TestPythonMLCallGraphShape {
   private static final TensorType TENSOR_2_NONE_INT32 =
       new TensorType(INT_32, asList(new NumericDim(2), null));
 
+  private static final TensorType TENSOR_2_NONE_FLOAT32 =
+      new TensorType(FLOAT_32, asList(new NumericDim(2), null));
+
+  private static final TensorType TENSOR_2_NONE_2_FLOAT32 =
+      new TensorType(FLOAT_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));
@@ -115,6 +121,15 @@ public class TestTensorflow2Model extends TestPythonMLCallGraphShape {
   private static final TensorType TENSOR_3_NONE_INT32 =
       new TensorType(INT_32, asList(new NumericDim(3), null));
 
+  private static final TensorType TENSOR_3_NONE_FLOAT32 =
+      new TensorType(FLOAT_32, asList(new NumericDim(3), null));
+
+  private static final TensorType TENSOR_3_NONE_NONE_FLOAT32 =
+      new TensorType(FLOAT_32, asList(new NumericDim(3), null, null));
+
+  private static final TensorType TENSOR_3_NONE_1_FLOAT32 =
+      new TensorType(FLOAT_32, asList(new NumericDim(3), null, new NumericDim(1)));
+
   private static final TensorType TENSOR_2_3_INT32 =
       new TensorType(INT_32, asList(new NumericDim(2), new NumericDim(3)));
 
@@ -4623,6 +4638,36 @@ public void testRaggedConstant6() throws ClassHierarchyException, CancelExceptio
     test("tf2_test_ragged_constant6.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_5_INT32)));
   }
 
+  @Test
+  public void testRaggedConstant7() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_ragged_constant7.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_3_NONE_INT32)));
+  }
+
+  @Test
+  public void testRaggedConstant8() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_ragged_constant8.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_3_NONE_FLOAT32)));
+  }
+
+  @Test
+  public void testRaggedConstant9() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_ragged_constant9.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_3_NONE_NONE_FLOAT32)));
+  }
+
+  @Test
+  public void testRaggedConstant10() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_ragged_constant10.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_3_NONE_1_FLOAT32)));
+  }
+
+  @Test
+  public void testRaggedConstant11() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_ragged_constant11.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_NONE_FLOAT32)));
+  }
+
+  @Test
+  public void testRaggedConstant12() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_ragged_constant12.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_NONE_2_FLOAT32)));
+  }
+
   private void test(
       String filename,
       String functionName,

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

@@ -1,5 +1,6 @@
 package com.ibm.wala.cast.python.ml.client;
 
+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;
@@ -28,7 +29,6 @@
 import java.util.ArrayList;
 import java.util.EnumSet;
 import java.util.List;
-import java.util.Optional;
 import java.util.Set;
 import java.util.logging.Logger;
 import java.util.stream.StreamSupport;
@@ -57,6 +57,65 @@ public RaggedConstant(PointsToSetVariable source) {
     super(source);
   }
 
+  private static Set<Integer> getPossibleInnerListLengths(
+      PropagationCallGraphBuilder builder, OrdinalSet<InstanceKey> pts) {
+    Set<Integer> ret = HashSetFactory.make();
+    PointerAnalysis<InstanceKey> pointerAnalysis = builder.getPointerAnalysis();
+
+    for (InstanceKey ik : pts) {
+      AllocationSiteInNode asin = getAllocationSiteInNode(ik);
+      TypeReference reference = asin.getConcreteType().getReference();
+
+      // A `list` or `tuple`.
+      if (reference.equals(list) || reference.equals(tuple)) {
+        OrdinalSet<InstanceKey> objectCatalogPointsToSet =
+            pointerAnalysis.getPointsToSet(
+                ((AstPointerKeyFactory) builder.getPointerKeyFactory())
+                    .getPointerKeyForObjectCatalog(asin));
+
+        assert objectCatalogPointsToSet.iterator().hasNext();
+
+        InstanceKey catalogIK =
+            objectCatalogPointsToSet
+                .iterator()
+                .next(); // Just need one element to check inner length.
+
+        ConstantKey<?> constantKey = (ConstantKey<?>) catalogIK;
+        Object constantKeyValue = constantKey.getValue();
+
+        Integer fieldIndex = (Integer) constantKeyValue;
+
+        FieldReference subscript =
+            FieldReference.findOrCreate(Root, findOrCreateAsciiAtom(fieldIndex.toString()), Root);
+
+        IField f = builder.getClassHierarchy().resolveField(subscript);
+
+        PointerKey pointerKeyForInstanceField = builder.getPointerKeyForInstanceField(asin, f);
+
+        OrdinalSet<InstanceKey> instanceFieldPointsToSet =
+            pointerAnalysis.getPointsToSet(pointerKeyForInstanceField);
+
+        boolean containsAllListsOrTuples =
+            StreamSupport.stream(instanceFieldPointsToSet.spliterator(), false)
+                .allMatch(
+                    ik -> {
+                      AllocationSiteInNode innerAsin = getAllocationSiteInNode(ik);
+
+                      if (innerAsin == null) return false;
+
+                      TypeReference innerReference = innerAsin.getConcreteType().getReference();
+                      return innerReference.equals(list) || innerReference.equals(tuple);
+                    });
+
+        if (!containsAllListsOrTuples) ret.add(objectCatalogPointsToSet.size());
+        else ret.addAll(getPossibleInnerListLengths(builder, instanceFieldPointsToSet));
+      } else
+        throw new IllegalStateException("Expected a list or tuple, but found: " + reference + ".");
+    }
+
+    return ret;
+  }
+
   private static Set<Integer> getPossibleOuterListLengths(
       PropagationCallGraphBuilder builder, OrdinalSet<InstanceKey> valuePointsToSet) {
     Set<Integer> ret = HashSetFactory.make();
@@ -246,31 +305,58 @@ protected Set<List<Dimension<?>>> getShapesOfValue(
       int K = maxDepth;
       LOGGER.fine("Tensor rank: " + K);
 
-      Optional<Integer> raggedRank = this.getRaggedRankArgumentValue(builder);
-      int R = raggedRank.orElse(K - 1);
-      LOGGER.fine("Ragged rank: " + R);
+      Set<Long> rankArguments = this.getPossibleRaggedRankArguments(builder);
 
-      // Step 3: Construct shape with rank K and ragged rank R.
+      if (rankArguments.isEmpty()) rankArguments.add(K - 1L); // Default ragged rank.
 
-      // Get the length of the outer list.
-      Set<Integer> possibleOuterListLengths =
-          getPossibleOuterListLengths(builder, valuePointsToSet);
+      for (Long R : rankArguments) {
+        LOGGER.fine("Ragged rank: " + R);
 
-      for (int outerListLength : possibleOuterListLengths) {
-        List<Dimension<?>> shape = new ArrayList<>();
-        shape.add(new NumericDim(outerListLength));
+        // Step 3: Construct shape with rank K and ragged rank R.
+        // The final shape is constructed by concatenating the Ragged Portion and the Uniform
+        // Portion.
 
-        // The first R dimensions are ragged.
-        for (int i = 0; i < R; i++) shape.add(null); // Unknown size for ragged dimensions.
+        // Part A: The Ragged Portion (Dimensions 0 to R)
 
-        /*
-        // The remaining K - R dimensions are dense.
-        for (int i = R; i < K; i++) {
-          shape.add(new NumericDim(-1)); // Unknown size for dense dimensions.
-        }
-        */
+        // For the ragged dimensions, TensorFlow does not look for a uniform length. It assigns the
+        // shape based on the row_splits.
 
-        ret.add(shape);
+        // Get the length of the outer list.
+        Set<Integer> possibleOuterListLengths =
+            getPossibleOuterListLengths(builder, valuePointsToSet);
+
+        for (int outerListLength : possibleOuterListLengths) {
+          List<Dimension<?>> shape = new ArrayList<>();
+
+          // Dim 0 (Batch): Always fixed. It is simply len(input_list).
+          shape.add(new NumericDim(outerListLength));
+
+          // The first R dimensions are ragged.
+          // Dim 1 to R: These are assigned None (or ? in older outputs) in the static shape,
+          // indicating they can vary.
+          for (Long i = 0L; i < R; i++) shape.add(null); // Unknown size for ragged dimensions.
+
+          // Part B: The Uniform Portion (Dimensions R + 1 to K)
+          // If R < K - 1 (meaning you requested fewer ragged dimensions than the total depth),
+          // TensorFlow enforces uniformity on the remaining inner dimensions.
+
+          // 1. It checks the length of every sub-list at these levels.
+          // 2. If any lengths differ, it throws a ValueError.
+          // 3. If they match, that length becomes the fixed size for that dimension.
+
+          if (R < K - 1) {
+            Set<Integer> possibleInnerListLengths =
+                getPossibleInnerListLengths(builder, valuePointsToSet);
+
+            // Determine the uniform lengths for dimensions R + 1 to K - 1.
+            for (long i = R + 1; i < K; i++) {
+              for (int innerListLength : possibleInnerListLengths)
+                shape.add(new NumericDim(innerListLength));
+            }
+          }
+
+          ret.add(shape);
+        }
       }
     }
 
@@ -288,9 +374,46 @@ private static int getMaximumDepthOfInstance(
       return 1 + getMaximumDepthOfEmptyList(builder, instance);
   }
 
-  private Optional<Integer> getRaggedRankArgumentValue(PropagationCallGraphBuilder builder) {
-    // TODO Auto-generated method stub
-    return Optional.empty();
+  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;
+  }
+
+  protected int getRaggedRankParameterPosition() {
+    return RAGGED_RANK.ordinal();
+  }
+
+  protected int getRaggedRankArgumentValueNumber(PropagationCallGraphBuilder builder) {
+    // TODO: Handle keyword arguments.
+    return this.getArgumentValueNumber(builder, this.getRaggedRankParameterPosition(), true);
   }
 
   /**
@@ -323,6 +446,7 @@ protected EnumSet<DType> getDefaultDTypes(PropagationCallGraphBuilder builder) {
       return EnumSet.of(FLOAT32);
     }
 
+    // Otherwise, there are values available to infer the dtype from.
     return super.getDefaultDTypes(builder);
   }
 }

com.ibm.wala.cast.python.test/data/tf2_test_ragged_constant10.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
+
+
+arg = [[[1], [2]], [[3]], [[4], [5], [6]]]
+
+x = tf.ragged.constant(arg, tf.float32, 1)
+assert x.shape == (3, None, 1)
+assert x.dtype == tf.float32
+
+f(x)

com.ibm.wala.cast.python.test/data/tf2_test_ragged_constant11.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
+
+
+arg = [[1, 2], [3]]
+
+x = tf.ragged.constant(arg, tf.float32, 1)
+
+assert x.shape == (2, None)
+assert x.dtype == tf.float32
+
+f(x)

com.ibm.wala.cast.python.test/data/tf2_test_ragged_constant12.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
+
+
+arg = [[[1, 2]], [[3, 4]]]
+
+x = tf.ragged.constant(arg, tf.float32, 1)
+
+assert x.shape == (2, None, 2)
+assert x.dtype == tf.float32
+
+f(x)

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

@@ -0,0 +1,24 @@
+# From https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/GradientTape#gradient.
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+arg = [[1, 2], [3], [4, 5, 6]]
+assert isinstance(arg, list)
+assert len(arg) == 3
+assert all(isinstance(row, list) for row in arg)
+assert all(isinstance(x, int) for row in arg for x in row)
+assert len(arg[0]) == 2
+assert len(arg[1]) == 1
+assert len(arg[2]) == 3
+
+x = tf.ragged.constant(arg, tf.int32)
+assert isinstance(x, tf.RaggedTensor)
+assert x.shape == (3, None)
+assert x.dtype == tf.int32
+
+f(x)

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

@@ -0,0 +1,24 @@
+# From https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/GradientTape#gradient.
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+arg = [[1, 2], [3], [4, 5, 6]]
+assert isinstance(arg, list)
+assert len(arg) == 3
+assert all(isinstance(row, list) for row in arg)
+assert all(isinstance(x, int) for row in arg for x in row)
+assert len(arg[0]) == 2
+assert len(arg[1]) == 1
+assert len(arg[2]) == 3
+
+x = tf.ragged.constant(arg, tf.float32)
+assert isinstance(x, tf.RaggedTensor)
+assert x.shape == (3, None)
+assert x.dtype == tf.float32
+
+f(x)