Progress on inner shapes.

Author: khatchad

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

@@ -120,6 +120,12 @@ public class TestTensorflow2Model extends TestPythonMLCallGraphShape {
   private static final TensorType TENSOR_2_NONE_2_INT32 =
       new TensorType(INT_32, asList(new NumericDim(2), null, new NumericDim(2)));
 
+  private static final TensorType TENSOR_2_NONE_2_3_INT32 =
+      new TensorType(INT_32, asList(new NumericDim(2), null, new NumericDim(2), new NumericDim(3)));
+
+  private static final TensorType TENSOR_2_NONE_2_2_INT32 =
+      new TensorType(INT_32, asList(new NumericDim(2), null, new NumericDim(2), new NumericDim(2)));
+
   @SuppressWarnings("unused")
   private static final TensorType TENSOR_2_NONE_NONE_NONE_INT32 =
       new TensorType(INT_32, asList(new NumericDim(2), null));
@@ -4697,6 +4703,22 @@ public void testRaggedConstant16() throws ClassHierarchyException, CancelExcepti
     test("tf2_test_ragged_constant16.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_NONE_2_INT32)));
   }
 
+  /**
+   * Test non-uniform inner dimensions.
+   *
+   * <p>FIXME: Should not throw an {@link AssertionError}.
+   */
+  @Test(expected = AssertionError.class)
+  public void testRaggedConstant17() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_ragged_constant17.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_NONE_2_3_INT32)));
+  }
+
+  /** This one works because the inner dimensions are uniform. */
+  @Test
+  public void testRaggedConstant18() throws ClassHierarchyException, CancelException, IOException {
+    test("tf2_test_ragged_constant18.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_NONE_2_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

@@ -380,6 +380,19 @@ protected Set<List<Dimension<?>>> getShapesOfValue(
             continue;
           }
 
+          if (!innerShapeArguments.isEmpty()) {
+            for (List<Dimension<?>> innerShape : innerShapeArguments) {
+              List<Dimension<?>> newShape = new ArrayList<>(shape.size());
+
+              newShape.addAll(shape);
+              newShape.addAll(innerShape);
+
+              ret.add(newShape);
+            }
+
+            continue;
+          }
+
           ret.add(shape);
         }
       }

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

@@ -0,0 +1,25 @@
+# From https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/GradientTape#gradient.
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+# A list of groups.
+# Group 1 has one 2x3 matrix.
+# Group 2 has two 2x3 matrices.
+data = [
+    [[[1, 1, 1], [2, 2, 2]]],  # Group 1
+    [[[3, 3, 3], [4, 4, 4]], [[5, 5, 5], [6, 6, 6]]],  # Group 2
+]
+
+# We set ragged_rank=1.
+# This means the OUTER list (the groups) varies in length.
+# But everything INSIDE a group is a fixed uniform block.
+t = tf.ragged.constant(data, None, 1)
+assert t.shape == (2, None, 2, 3)
+assert t.dtype == tf.int32
+
+f(t)

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

@@ -0,0 +1,19 @@
+# From https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/GradientTape#gradient.
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+# Two groups. Group 1 has one matrix. Group 2 has two matrices.
+data = [[[[1, 1], [2, 2]]], [[[3, 3], [4, 4]], [[5, 5], [6, 6]]]]  # Group 1  # Group 2
+
+# We set ragged_rank=1 (The groups are ragged).
+# We set inner_shape=(2, 2) (The things inside are 2x2 matrices).
+t = tf.ragged.constant(data, None, 1, (2, 2))
+assert t.shape == (2, None, 2, 2)
+assert t.dtype == tf.int32
+
+f(t)