More tests.

Author: khatchad

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

@@ -77,6 +77,12 @@ public class TestTensorflow2Model extends TestPythonMLCallGraphShape {
   private static final TensorType TENSOR_2_1_FLOAT32 =
       new TensorType(FLOAT_32, asList(new NumericDim(2), new NumericDim(1)));
 
+  private static final TensorType TENSOR_2_3_3_FLOAT32 =
+      new TensorType(FLOAT_32, asList(new NumericDim(2), new NumericDim(3), new NumericDim(3)));
+
+  private static final TensorType TENSOR_2_3_4_FLOAT32 =
+      new TensorType(FLOAT_32, asList(new NumericDim(2), new NumericDim(3), new NumericDim(4)));
+
   private static final TensorType TENSOR_2_FLOAT32 =
       new TensorType(FLOAT_32, asList(new NumericDim(2)));
 
@@ -183,6 +189,24 @@ public void testFunction8()
         Map.of(2, Set.of(TENSOR_2_1_FLOAT32, TENSOR_2_FLOAT32)));
   }
 
+  @Test
+  public void testFunction9()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_function9.py", "func", 1, 1, Map.of(2, Set.of(TENSOR_1_2_FLOAT32)));
+  }
+
+  @Test
+  public void testFunction10()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_function10.py", "func", 1, 1, Map.of(2, Set.of(TENSOR_2_3_4_FLOAT32)));
+  }
+
+  @Test
+  public void testFunction11()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_function11.py", "func", 1, 1, Map.of(2, Set.of(TENSOR_2_3_3_FLOAT32)));
+  }
+
   @Test
   public void testDecorator()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {

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

@@ -0,0 +1,16 @@
+import tensorflow as tf
+
+
+def func(t):
+    pass
+
+
+a = tf.constant(
+    [
+        [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
+        [[13, 14, 15, 16], [17, 18, 19, 20], [21, 22, 23, 24]],
+    ]
+)
+assert a.shape == (2, 3, 4)
+
+func(a)

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

@@ -0,0 +1,16 @@
+import tensorflow as tf
+
+
+def func(t):
+    pass
+
+
+a = tf.constant(
+    [
+        [[1, 2, 3], [5, 6, 7], [9, 10, 11]],
+        [[13, 14, 15], [17, 18, 19], [21, 22, 23]],
+    ]
+)
+assert a.shape == (2, 3, 3)
+
+func(a)

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

@@ -0,0 +1,11 @@
+import tensorflow as tf
+
+
+def func(t):
+    pass
+
+
+a = tf.constant([[1.0, 3.0]])
+assert a.shape == (1, 2)
+
+func(a)