Progress.

Author: khatchad

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

@@ -4692,6 +4692,11 @@ public void testRaggedConstant15() throws ClassHierarchyException, CancelExcepti
     test("tf2_test_ragged_constant15.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_NONE_2_INT32)));
   }
 
+  @Test
+  public void testRaggedConstant16() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_ragged_constant16.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

@@ -32,6 +32,7 @@
 import java.util.ArrayList;
 import java.util.EnumSet;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 import java.util.logging.Logger;
 import java.util.stream.StreamSupport;
@@ -308,8 +309,15 @@ protected Set<List<Dimension<?>>> getShapesOfValue(
       int K = maxDepth;
       LOGGER.fine("Tensor rank: " + K);
 
-      Set<Long> rankArguments = this.getPossibleRaggedRankArguments(builder);
-      Set<List<Dimension<?>>> innerShapeArguments = this.getPossibleInnerShapeArguments(builder);
+      Set<Long> rankArguments =
+          this.getPossibleRaggedRankArguments(builder).stream()
+              .filter(Objects::nonNull)
+              .collect(toSet());
+
+      Set<List<Dimension<?>>> innerShapeArguments =
+          this.getPossibleInnerShapeArguments(builder).stream()
+              .filter(Objects::nonNull)
+              .collect(toSet());
 
       if (rankArguments.isEmpty())
         // Default ragged rank.

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

@@ -738,6 +738,15 @@ protected Set<Integer> getNumberOfPossiblePositionalArguments(
     return ret;
   }
 
+  /**
+   * Returns the possible long arguments for the given value number. If the argument is `None`, then
+   * a null value will be contained within the returned set.
+   *
+   * @param builder The {@link PropagationCallGraphBuilder} used for the analysis.
+   * @param valueNumber The value number of the argument.
+   * @return A set of possible long arguments. If the argument is `None`, then a null value will be
+   *     contained within the returned set.
+   */
   protected Set<Long> getPossibleLongArguments(
       PropagationCallGraphBuilder builder, int valueNumber) {
     Set<Long> ret = HashSetFactory.make();
@@ -753,14 +762,17 @@ protected Set<Long> getPossibleLongArguments(
             "Empty points-to set in source: " + this.getSource() + ".");
 
       for (InstanceKey instanceKey : pointsToSet)
-        if (instanceKey instanceof com.ibm.wala.ipa.callgraph.propagation.ConstantKey) {
+        if (instanceKey instanceof ConstantKey) {
           ConstantKey<?> constantKey = (ConstantKey<?>) instanceKey;
           Object constantKeyValue = constantKey.getValue();
 
           if (constantKeyValue instanceof Long) {
             Long value = (Long) constantKeyValue;
             ret.add(value);
-          } else
+          } else if (constantKeyValue == null)
+            // The argument may be `None`.
+            ret.add(null);
+          else
             throw new IllegalStateException(
                 "Expected a long, but found: " + constantKeyValue + ".");
         } else

com.ibm.wala.cast.python.test/data/tf2_test_ragged_constant16.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, None, (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)