Initial work on dtype argument for tf.range().

Author: khatchad

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

@@ -2036,6 +2036,12 @@ public void testTFRange3()
     test("tf2_test_tf_range3.py", "f", 1, 1, Map.of(2, Set.of(SCALAR_TENSOR_OF_INT32)));
   }
 
+  @Test
+  public void testTFRange4()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_tf_range4.py", "f", 1, 1, Map.of(2, Set.of(SCALAR_TENSOR_OF_FLOAT32)));
+  }
+
   @Test
   public void testImport()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {

com.ibm.wala.cast.python.ml/data/tensorflow.xml

@@ -537,7 +537,7 @@
         <method name="read_data" descriptor="()LRoot;">
           <new def="x" class="Llist" />
           <new def="z" class="Ltensorflow/functions/constant" />
-          <constant name="p" type="int" value="1" />
+          <constant name="p" type="int" value="1" /> <!-- FIXME: Unfortunately, if a dtype argument is supplied to `tf.range()`, the value may not be an integer. -->
           <call class="Ltensorflow/functions/constant" name="do" descriptor="()LRoot;" type="virtual" arg0="z" arg1="p" numArgs="2" def="y" />
           <putfield class="LRoot" field="0" fieldType="LRoot" ref="x" value="y" />
           <return value="x" />

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

@@ -43,6 +43,8 @@ public class Range extends TensorGenerator {
 
   private static final String FUNCTION_NAME = "tf.range()";
 
+  private static final int DTYPE_POSITIONAL_PARAMETER_INDEX = 3;
+
   public Range(PointsToSetVariable source, CGNode node) {
     super(source, node);
   }
@@ -232,10 +234,13 @@ protected int getValueNumberForShapeArgument() {
 
   @Override
   protected int getValueNumberForDTypeArgument() {
+    // TODO: Handle keyword arguments.
+
     // TODO: We need a value number for the dtype argument. Also, that value number can differ
     // depending on the version of the `range` function being called.
-
-    return -1; // Positional dtype argument for range() is not yet implemented.
+    return this.getNode().getIR().getMethod().isStatic()
+        ? this.getNode().getIR().getParameter(DTYPE_POSITIONAL_PARAMETER_INDEX)
+        : this.getNode().getIR().getParameter(DTYPE_POSITIONAL_PARAMETER_INDEX + 1);
   }
 
   @Override

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

@@ -0,0 +1,25 @@
+# From: https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/range#for_example
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+start = 3
+limit = 18
+delta = 3
+
+r = tf.range(start, limit, delta, tf.float32)
+assert isinstance(r, tf.Tensor)
+assert r.shape == (5,)
+assert r.dtype == tf.float32
+
+for i in r:
+    assert isinstance(i, tf.Tensor)
+    assert (
+        i.dtype == tf.float32
+    )  # NOTE: This is getting cast here from the original input.
+    assert i.shape == ()
+    f(i)