Add test.

Author: khatchad

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

@@ -4390,6 +4390,16 @@ public void testEye5()
     test("tf2_test_eye5.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_3_2_2_3_FLOAT32)));
   }
 
+  /**
+   * FIXME: Should not throw an {@link IllegalArgumentException} once
+   * https://github.com/wala/ML/issues/340 is fixed.
+   */
+  @Test(expected = IllegalArgumentException.class)
+  public void testEye6()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_eye6.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_3_2_2_3_FLOAT32)));
+  }
+
   private void test(
       String filename,
       String functionName,

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

@@ -0,0 +1,17 @@
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+batch = tf.constant([3, 2])
+assert batch.shape == (2,)
+assert batch.dtype == tf.int32
+
+arg = tf.eye(2, 3, batch)
+
+assert arg.shape == (3, 2, 2, 3)
+assert arg.dtype == tf.float32
+
+f(arg)