Add future test.

Author: khatchad

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

@@ -95,6 +95,9 @@ public class TestTensorflow2Model extends TestPythonMLCallGraphShape {
   private static final TensorType TENSOR_3_3_INT32 =
       new TensorType(INT_32, asList(new NumericDim(3), new NumericDim(3)));
 
+  private static final TensorType TENSOR_2_3_INT32 =
+      new TensorType(INT_32, asList(new NumericDim(2), new NumericDim(3)));
+
   private static final TensorType TENSOR_2_1_FLOAT32 =
       new TensorType(FLOAT_32, asList(new NumericDim(2), new NumericDim(1)));
 
@@ -104,6 +107,9 @@ public class TestTensorflow2Model extends TestPythonMLCallGraphShape {
   private static final TensorType TENSOR_2_3_4_INT32 =
       new TensorType(INT_32, asList(new NumericDim(2), new NumericDim(3), new NumericDim(4)));
 
+  private static final TensorType TENSOR_2_5_3_FLOAT32 =
+      new TensorType(FLOAT_32, asList(new NumericDim(2), new NumericDim(5), new NumericDim(3)));
+
   private static final TensorType TENSOR_20_28_28_FLOAT32 =
       new TensorType(FLOAT_32, asList(new NumericDim(20), new NumericDim(28), new NumericDim(28)));
 
@@ -4314,6 +4320,23 @@ public void testOneHot17()
     test("tf2_test_one_hot17.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_3_FLOAT32)));
   }
 
+  @Test
+  public void testOneHot18()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_one_hot18.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_5_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 testOneHot19()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_one_hot19.py", "g", 1, 1, Map.of(2, Set.of(TENSOR_2_3_INT32)));
+    test("tf2_test_one_hot19.py", "f", 1, 1, Map.of(2, Set.of(TENSOR_2_5_3_FLOAT32)));
+  }
+
   private void test(
       String filename,
       String functionName,

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

@@ -0,0 +1,21 @@
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+arg1 = [[10, 20, 30], [40, 50, 60]]  # Row 1  # Row 2
+
+assert isinstance(arg1, list)
+assert all(isinstance(row, list) for row in arg1)
+assert all(isinstance(elem, int) for row in arg1 for elem in row)
+assert len(arg1) == 2
+assert all(len(row) == 3 for row in arg1)
+
+arg2 = tf.one_hot(arg1, 5, None, None, 1)
+assert isinstance(arg2, tf.Tensor)
+assert arg2.dtype == tf.float32
+assert arg2.shape == (2, 5, 3)
+
+f(arg2)

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

@@ -0,0 +1,24 @@
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+def g(a):
+    pass
+
+
+my_tensor = tf.constant([[10, 20, 30], [40, 50, 60]])  # Row 1  # Row 2
+assert isinstance(my_tensor, tf.Tensor)
+assert my_tensor.dtype == tf.int32
+assert my_tensor.shape == (2, 3)
+
+g(my_tensor)
+
+arg = tf.one_hot(my_tensor, 5, None, None, 1)
+assert isinstance(arg, tf.Tensor)
+assert arg.dtype == tf.float32
+assert arg.shape == (2, 5, 3)
+
+f(arg)