Handle multiple parameters for tf.range().

Author: khatchad

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

@@ -1328,13 +1328,23 @@ public void testAdd37()
   @Test
   public void testAdd38()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
-    test("tf2_test_add38.py", "add", 2, 2, Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
+    test(
+        "tf2_test_add38.py",
+        "add",
+        2,
+        2,
+        Map.of(2, Set.of(TENSOR_5_INT32), 3, Set.of(TENSOR_5_INT32)));
   }
 
   @Test
   public void testAdd39()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
-    test("tf2_test_add39.py", "add", 2, 2, Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
+    test(
+        "tf2_test_add39.py",
+        "add",
+        2,
+        2,
+        Map.of(2, Set.of(TENSOR_5_INT32), 3, Set.of(TENSOR_5_INT32)));
   }
 
   @Test
@@ -1422,7 +1432,12 @@ public void testAdd49()
   @Test
   public void testAdd50()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
-    test("tf2_test_add50.py", "add", 2, 2, Map.of(2, Set.of(MNIST_INPUT), 3, Set.of(MNIST_INPUT)));
+    test(
+        "tf2_test_add50.py",
+        "add",
+        2,
+        2,
+        Map.of(2, Set.of(TENSOR_5_INT32), 3, Set.of(TENSOR_5_INT32)));
   }
 
   @Test

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

@@ -18,7 +18,6 @@
 import com.ibm.wala.ipa.callgraph.propagation.cfa.CallString;
 import com.ibm.wala.ssa.SSAAbstractInvokeInstruction;
 import com.ibm.wala.util.collections.HashSetFactory;
-import com.ibm.wala.util.debug.UnimplementedError;
 import com.ibm.wala.util.intset.OrdinalSet;
 import java.util.EnumSet;
 import java.util.Iterator;
@@ -64,10 +63,10 @@ protected Set<List<Dimension<?>>> getShapes(PropagationCallGraphBuilder builder)
     // 2. `tf.range(start, limit, delta)` - generates a range from start to limit with a step of
     // delta.
 
-    // First, decide which version of the `range` function is being called based on the number of
-    // numeric arguments.j
+    // Decide which version of the `range` function is being called based on the number of numeric
+    // arguments.
     // TODO: Handle keyword arguments.
-    for (Integer numOfPoisitionArguments : getNumberOfPossiblePositionalArguments(builder)) {
+    for (Integer numOfPoisitionArguments : getNumberOfPossiblePositionalArguments(builder))
       if (numOfPoisitionArguments == 1) {
         // it must *just* be `limit`.
         PointerKey limitPK =
@@ -81,11 +80,42 @@ protected Set<List<Dimension<?>>> getShapes(PropagationCallGraphBuilder builder)
           int shape = (int) Math.ceil((limit - start) / delta);
           ret.add(List.of(new NumericDim(shape))); // Add the shape as a 1D tensor.
         }
+      } else if (numOfPoisitionArguments == 3) {
+        // it must be `start`, `limit`, and `delta`.
+        PointerKey startPK =
+            pointerAnalysis.getHeapModel().getPointerKeyForLocal(this.getNode(), 2);
+        PointerKey limitPK =
+            pointerAnalysis.getHeapModel().getPointerKeyForLocal(this.getNode(), 3);
+        PointerKey deltaPK =
+            pointerAnalysis.getHeapModel().getPointerKeyForLocal(this.getNode(), 4);
+
+        OrdinalSet<InstanceKey> startPointsToSet = pointerAnalysis.getPointsToSet(startPK);
+        OrdinalSet<InstanceKey> limitPointsToSet = pointerAnalysis.getPointsToSet(limitPK);
+        OrdinalSet<InstanceKey> deltaPointsToSet = pointerAnalysis.getPointsToSet(deltaPK);
+
+        assert !startPointsToSet.isEmpty() : "Expected a non-empty points-to set for start.";
+        assert !limitPointsToSet.isEmpty() : "Expected a non-empty points-to set for limit.";
+        assert !deltaPointsToSet.isEmpty() : "Expected a non-empty points-to set for delta.";
+
+        for (InstanceKey startIK : startPointsToSet) {
+          start = ((Number) ((ConstantKey<?>) startIK).getValue()).doubleValue();
+
+          for (InstanceKey limitIK : limitPointsToSet) {
+            limit = ((Number) ((ConstantKey<?>) limitIK).getValue()).doubleValue();
+
+            for (InstanceKey deltaIK : deltaPointsToSet) {
+              delta = ((Number) ((ConstantKey<?>) deltaIK).getValue()).doubleValue();
+
+              int shape = (int) Math.ceil((limit - start) / delta);
+              ret.add(List.of(new NumericDim(shape))); // Add the shape as a 1D tensor.
+            }
+          }
+        }
       } else
-        // TODO: Handle more cases.
-        throw new UnimplementedError(
-            "Currently cannot handle more than one numeric positional argument for range().");
-    }
+        throw new IllegalStateException(
+            "Expected either 1 or 3 positional arguments for range(), but got: "
+                + numOfPoisitionArguments
+                + ".");
 
     return ret;
   }

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

@@ -5,4 +5,12 @@ def add(a, b):
     return a + b
 
 
-c = add(tf.range(3, 18, 3), tf.range(5))
+arg1 = tf.range(3, 18, 3)
+assert arg1.shape == (5,)
+assert arg1.dtype == tf.int32
+
+arg2 = tf.range(5)
+assert arg2.shape == (5,)
+assert arg2.dtype == tf.int32
+
+c = add(arg1, arg2)

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

@@ -12,6 +12,12 @@ def f(a):
 delta = 3
 
 r = tf.range(start, limit, delta)
+assert isinstance(r, tf.Tensor)
+assert r.shape == (5,)
+assert r.dtype == tf.int32
 
 for i in r:
+    assert isinstance(i, tf.Tensor)
+    assert i.dtype == tf.int32
+    assert i.shape == ()
     f(i)