ponder-lab · khatchad · Jul 11, 2025 · Jul 14, 2025 · Jul 14, 2025 · Jul 16, 2025
@@ -154,6 +154,9 @@
         <putfield class="LRoot" field="sparse_tensor" fieldType="LRoot" ref="framework" value="sparse_tensor" />
         <new def="constant_op" class="Lobject" />
         <putfield class="LRoot" field="constant_op" fieldType="LRoot" ref="framework" value="constant_op" />
+        <!-- https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/dtypes -->
+        <new def="dtypes" class="Lobject" />
+        <putfield class="LRoot" field="dtypes" fieldType="LRoot" ref="framework" value="dtypes" />
         <new def="ragged_factory_ops" class="Lobject" />
         <putfield class="LRoot" field="ragged_factory_ops" fieldType="LRoot" ref="ragged" value="ragged_factory_ops" />
         <new def="ragged_math_ops" class="Lobject" />
@@ -218,6 +221,10 @@
         <new def="constant" class="Ltensorflow/functions/constant" />
         <putfield class="LRoot" field="constant" fieldType="LRoot" ref="x" value="constant" />
         <putfield class="LRoot" field="constant" fieldType="LRoot" ref="constant_op" value="constant" />
+        <!-- https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/dtypes#float32 -->
+        <new def="float32" class="Ltensorflow/dtypes/DType" />
+        <putfield class="LRoot" field="float32" fieldType="LRoot" ref="x" value="float32" />
+        <putfield class="LRoot" field="float32" fieldType="LRoot" ref="dtypes" value="float32" />
         <new def="ragged_constant" class="Ltensorflow/functions/ragged_constant" />
         <putfield class="LRoot" field="constant" fieldType="LRoot" ref="ragged" value="ragged_constant" />
         <putfield class="LRoot" field="constant" fieldType="LRoot" ref="ragged_factory_ops" value="ragged_constant" />
@@ -315,6 +322,17 @@
         </method>
       </class>
     </package>
+    <package name="tensorflow/dtypes">
+      <class name="DType" allocatable="true">
+        <!-- https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/dtypes/DType -->
+        <method name="do" descriptor="()LRoot;" numArgs="3" paramNames="self type_enum handle_data">
+          <new def="obj" class="Ltensorflow//DType" />
+          <putfield class="LRoot" field="type_enum" fieldType="LRoot" ref="obj" value="type_enum" />
+          <putfield class="LRoot" field="handle_data" fieldType="LRoot" ref="obj" value="handle_data" />
+          <return value="obj" />
+        </method>
+      </class>
+    </package>
     <package name="tensorflow/objects">
       <class name="feature" allocatable="true" />
     </package>

@@ -447,11 +447,11 @@ public String toString() {
     };
   }
 
-  private final Map<PointsToSetVariable, TensorType> init;
+  private final Map<PointsToSetVariable, Set<TensorType>> init;
 
   public TensorTypeAnalysis(
       Graph<PointsToSetVariable> G,
-      Map<PointsToSetVariable, TensorType> init,
+      Map<PointsToSetVariable, Set<TensorType>> init,
       Map<PointsToSetVariable, TensorType> reshapeTypes,
       Map<PointsToSetVariable, TensorType> set_shapes,
       Set<PointsToSetVariable> conv2ds,
@@ -480,7 +480,8 @@ protected TensorVariable[] makeStmtRHS(int size) {
   protected void initializeVariables() {
     super.initializeVariables();
     for (PointsToSetVariable src : init.keySet()) {
-      getOut(src).state.add(init.get(src));
+      Set<TensorType> tensorTypes = init.get(src);
+      getOut(src).state.addAll(tensorTypes);
     }
   }
 

@@ -0,0 +1,60 @@
+package com.ibm.wala.cast.python.ml.client;
+
+import com.ibm.wala.cast.python.ml.types.TensorFlowTypes.DType;
+import com.ibm.wala.cast.python.ml.types.TensorType.Dimension;
+import com.ibm.wala.ipa.callgraph.CGNode;
+import com.ibm.wala.ipa.callgraph.propagation.PointsToSetVariable;
+import com.ibm.wala.ipa.callgraph.propagation.PropagationCallGraphBuilder;
+import java.util.EnumSet;
+import java.util.List;
+import java.util.Set;
+
+/**
+ * Represents a call to the <code>constant()</code> function in TensorFlow.
+ *
+ * @see <a href="https://www.tensorflow.org/api_docs/python/tf/constant">constant()</a>.
+ * @author <a href="mailto:khatchad@hunter.cuny.edu">Raffi Khatchadourian</a>
+ */
+public class Constant extends TensorGenerator {
+
+  private static final int VALUE_NUMBER_FOR_VALUE_ARGUMENT = 2;
+
+  private static final int VALUE_NUMBER_FOR_DTYPE_ARGUMENT = 3;
+
+  private static final int VALUE_NUMBER_FOR_SHAPE_ARGUMENT = 4;
+
+  public Constant(PointsToSetVariable source, CGNode node) {
+    super(source, node);
+  }
+
+  @Override
+  protected Set<List<Dimension<?>>> getDefaultShapes(PropagationCallGraphBuilder builder) {
+    // If the shape argument is not specified, then the shape is inferred from the shape of value.
+    // TODO: Handle keyword arguments.
+    return getShapes(builder, this.getValueNumberForValueArgument());
+  }
+
+  @Override
+  protected EnumSet<DType> getDefaultDTypes(PropagationCallGraphBuilder builder) {
+    // If the dtype argument is not specified, then the type is inferred from the type of value.
+    // TODO: Handle keyword arguments.
+    return getDTypes(builder, this.getValueNumberForValueArgument());
+  }
+
+  @Override
+  protected int getValueNumberForDTypeArgument() {
+    return VALUE_NUMBER_FOR_DTYPE_ARGUMENT;
+  }
+
+  protected int getValueNumberForValueArgument() {
+    return VALUE_NUMBER_FOR_VALUE_ARGUMENT;
+  }
+
+  @Override
+  protected int getValueNumberForShapeArgument() {
+    // Shapes can also be specified as an explicit argument. Here, we examine the third explicit
+    // argument (recall that the first argument is implicit and corresponds to the called
+    // function's name).
+    return VALUE_NUMBER_FOR_SHAPE_ARGUMENT;
+  }
+}
@@ -0,0 +1,54 @@
+package com.ibm.wala.cast.python.ml.client;
+
+import static com.ibm.wala.cast.python.ml.types.TensorFlowTypes.DType.FLOAT32;
+
+import com.ibm.wala.cast.python.ml.types.TensorFlowTypes.DType;
+import com.ibm.wala.cast.python.ml.types.TensorType.Dimension;
+import com.ibm.wala.ipa.callgraph.CGNode;
+import com.ibm.wala.ipa.callgraph.propagation.PointsToSetVariable;
+import com.ibm.wala.ipa.callgraph.propagation.PropagationCallGraphBuilder;
+import java.util.EnumSet;
+import java.util.List;
+import java.util.Set;
+
+/**
+ * A generator for tensors created by the `ones()` function in TensorFlow.
+ *
+ * @see <a href="https://www.tensorflow.org/api_docs/python/tf/ones">TensorFlow ones() API</a>.
+ * @author <a href="mailto:khatchad@hunter.cuny.edu">Raffi Khatchadourian</a>
+ */
+public class Ones extends TensorGenerator {
+
+  private static final int VALUE_NUMBER_FOR_SHAPE_ARGUMENT = 2;
+
+  private static final int VALUE_NUMBER_FOR_DTYPE_ARGUMENT = 3;
+
+  public Ones(PointsToSetVariable source, CGNode node) {
+    super(source, node);
+  }
+
+  @Override
+  protected EnumSet<DType> getDefaultDTypes(PropagationCallGraphBuilder builder) {
+    LOGGER.info(
+        "No dtype specified for source: " + source + ". Using default dtype of: " + FLOAT32 + " .");
+
+    // Use the default dtype of float32.
+    return EnumSet.of(FLOAT32);
+  }
+
+  @Override
+  protected Set<List<Dimension<?>>> getDefaultShapes(PropagationCallGraphBuilder builder) {
+    throw new UnsupportedOperationException(
+        "Shapes for ones() are mandatory and must be provided explicitly.");
+  }
+
+  @Override
+  protected int getValueNumberForShapeArgument() {
+    return VALUE_NUMBER_FOR_SHAPE_ARGUMENT; // The shape is in the first explicit argument.
+  }
+
+  @Override
+  protected int getValueNumberForDTypeArgument() {
+    return VALUE_NUMBER_FOR_DTYPE_ARGUMENT; // The dtype is in the second explicit argument.
+  }
+}
@@ -667,10 +667,9 @@ public TensorTypeAnalysis performAnalysis(PropagationCallGraphBuilder builder)
     Set<PointsToSetVariable> sources =
         getDataflowSources(dataflow, builder.getCallGraph(), builder.getPointerAnalysis());
 
-    TensorType mnistData = TensorType.mnistInput();
-    Map<PointsToSetVariable, TensorType> init = HashMapFactory.make();
+    Map<PointsToSetVariable, Set<TensorType>> init = HashMapFactory.make();
 
-    for (PointsToSetVariable v : sources) init.put(v, mnistData);
+    for (PointsToSetVariable v : sources) init.put(v, getTensorTypes(v, builder));
 
     Map<PointsToSetVariable, TensorType> placeholders = null;
     try {
@@ -681,7 +680,7 @@ public TensorTypeAnalysis performAnalysis(PropagationCallGraphBuilder builder)
     logger.fine("Placeholders: " + placeholders);
 
     for (Map.Entry<PointsToSetVariable, TensorType> e : placeholders.entrySet())
-      init.put(e.getKey(), e.getValue());
+      init.put(e.getKey(), Set.of(e.getValue()));
 
     Map<PointsToSetVariable, TensorType> setCalls = HashMapFactory.make();
     Map<PointsToSetVariable, TensorType> set_shapes = getShapeSourceCalls(set_shape, builder, 1);
@@ -722,6 +721,24 @@ public TensorTypeAnalysis performAnalysis(PropagationCallGraphBuilder builder)
     return tt;
   }
 
+  /**
+   * Returns the set of possible {@link TensorType}s that the given {@link PointsToSetVariable} can
+   * take on.
+   *
+   * @param source The dataflow source to analyze.
+   * @param builder The {@link PropagationCallGraphBuilder} used to build the call graph and pointer
+   *     analysis.
+   * @return A set of {@link TensorType}s that the given {@link PointsToSetVariable} can take on.
+   *     Empty set is returned if the possible tensor types cannot be determined.
+   */
+  private Set<TensorType> getTensorTypes(
+      PointsToSetVariable source, PropagationCallGraphBuilder builder) {
+    logger.info("Getting tensor types for source: " + source + ".");
+
+    TensorGenerator generator = TensorGeneratorFactory.getGenerator(source);
+    return generator.getTensorTypes(builder);
+  }
+
   private Map<PointsToSetVariable, TensorType> handleShapeSourceOp(
       PropagationCallGraphBuilder builder,
       Graph<PointsToSetVariable> dataflow,

@@ -0,0 +1,166 @@
+package com.ibm.wala.cast.python.ml.client;
+
+import static java.util.function.Function.identity;
+
+import com.ibm.wala.cast.python.ml.types.TensorFlowTypes.DType;
+import com.ibm.wala.cast.python.ml.types.TensorType.Dimension;
+import com.ibm.wala.cast.python.ml.types.TensorType.NumericDim;
+import com.ibm.wala.ipa.callgraph.CGNode;
+import com.ibm.wala.ipa.callgraph.propagation.ConstantKey;
+import com.ibm.wala.ipa.callgraph.propagation.InstanceKey;
+import com.ibm.wala.ipa.callgraph.propagation.PointerAnalysis;
+import com.ibm.wala.ipa.callgraph.propagation.PointerKey;
+import com.ibm.wala.ipa.callgraph.propagation.PointsToSetVariable;
+import com.ibm.wala.ipa.callgraph.propagation.PropagationCallGraphBuilder;
+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.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.StreamSupport;
+
+/**
+ * A representation of the TensorFlow range operation.
+ *
+ * <p>This class is used to generate a tensor that contains a sequence of numbers, similar to the
+ * range function in Python.
+ *
+ * @see <a href="https://www.tensorflow.org/api_docs/python/tf/range">TensorFlow range
+ *     documentation</a>.
+ * @author <a href="mailto:khatchad@hunter.cuny.edu">Raffi Khatchadourian</a>
+ */
+public class Range extends TensorGenerator {
+
+  public Range(PointsToSetVariable source, CGNode node) {
+    super(source, node);
+  }
+
+  @Override
+  protected Set<List<Dimension<?>>> getShapes(PropagationCallGraphBuilder builder) {
+    Set<List<Dimension<?>>> ret = HashSetFactory.make();
+    PointerAnalysis<InstanceKey> pointerAnalysis = builder.getPointerAnalysis();
+
+    // The shape of a range tensor is always a 1D tensor with the length equal to the number of
+    // elements in the range.
+    // For example, `tf.range(5)` produces a tensor with shape (5,).
+
+    double start = 0; // Default start value.
+    double limit = start; // Default limit value.
+    double delta = 1; // Default step value.
+
+    // There are two versions of the `range` function:
+    // 1. `tf.range(limit)` - generates a range from 0 to limit
+    // 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
+    // TODO: Handle keyword arguments.
+
+    int numOfNumericPositionalArgs = getNumberOfNumericPositionalArgs(pointerAnalysis);
+
+    if (numOfNumericPositionalArgs == 1) {
+      // it must *just* be `limit`.
+      PointerKey limitPK = pointerAnalysis.getHeapModel().getPointerKeyForLocal(node, 2);
+      OrdinalSet<InstanceKey> limitPointsToSet = pointerAnalysis.getPointsToSet(limitPK);
+
+      assert !limitPointsToSet.isEmpty() : "Expected a non-empty points-to set for limit.";
+
+      for (InstanceKey limitIK : limitPointsToSet)
+        if (limitIK instanceof ConstantKey) {
+          limit = ((Number) ((ConstantKey<?>) limitIK).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
+          throw new IllegalStateException(
+              "Expected a " + ConstantKey.class + " for limit, but got: " + limitIK + ".");
+    } else
+      // TODO: Handle more cases.
+      throw new UnimplementedError(
+          "Currently cannot handle more than one numeric positional argument for range().");
+
+    return ret;
+  }
+
+  private int getNumberOfNumericPositionalArgs(PointerAnalysis<InstanceKey> pointerAnalysis) {
+    int ret = 0;
+    int explicitArgumentIndex = 2; // Start from the first explicit argument.
+
+    while (true) {
+      PointerKey pk =
+          pointerAnalysis.getHeapModel().getPointerKeyForLocal(node, explicitArgumentIndex);
+      OrdinalSet<InstanceKey> pointsToSet = pointerAnalysis.getPointsToSet(pk);
+
+      if (pointsToSet.isEmpty()) break; // End of positional arguments.
+
+      // Check if the pointsToSet contains numeric values.
+      boolean allNumeric =
+          StreamSupport.stream(pointsToSet.spliterator(), false)
+              .filter(ik -> ik instanceof ConstantKey)
+              .map(ik -> (ConstantKey<?>) ik)
+              .map(ConstantKey::getValue)
+              .allMatch(v -> v instanceof Number); // Check if all values are numeric.
+
+      if (!allNumeric) break; // There's some argument that is not numeric for this argument.
+
+      ret++; // Increment the count of numeric positional arguments.
+      explicitArgumentIndex++; // Move to the next explicit argument.
+    }
+
+    return ret;
+  }
+
+  @Override
+  protected EnumSet<DType> getDefaultDTypes(PropagationCallGraphBuilder builder) {
+    // The dtype of the resulting tensor is inferred from the inputs unless it is provided
+    // explicitly.
+
+    // TODO: Handle keyword arguments.
+    int numberOfNumericPositionalArgs =
+        getNumberOfNumericPositionalArgs(builder.getPointerAnalysis());
+
+    EnumSet<DType> types =
+        IntStream.range(0, numberOfNumericPositionalArgs)
+            .map(i -> i + 2) // Positional arguments start at index 2.
+            .mapToObj(val -> getDTypes(builder, val).stream())
+            .flatMap(identity())
+            .distinct()
+            .collect(Collectors.toCollection(() -> EnumSet.noneOf(DType.class)));
+
+    // FIXME: We can't tell the difference here between varying dtypes in a single call and that of
+    // possible varying dtypes values from the points-to graph. Below, we are treating it as these
+    // values lie in a single call, but that may not be the case.
+
+    if (types.contains(DType.FLOAT64)) return EnumSet.of(DType.FLOAT64);
+    else if (types.contains(DType.FLOAT32)) return EnumSet.of(DType.FLOAT32);
+    else if (types.contains(DType.INT64)) return EnumSet.of(DType.INT64);
+    else if (types.contains(DType.INT32)) return EnumSet.of(DType.INT32);
+
+    throw new IllegalStateException(
+        "Expected at least one numeric dtype for range(), but got: " + types + ".");
+  }
+
+  @Override
+  protected Set<List<Dimension<?>>> getDefaultShapes(PropagationCallGraphBuilder builder) {
+    throw new UnsupportedOperationException(
+        "Shapes for range() are derived from mandatory numeric arguments and must be provided"
+            + " explicitly.");
+  }
+
+  @Override
+  protected int getValueNumberForShapeArgument() {
+    throw new UnsupportedOperationException(
+        "Range does not have a shape argument. Its shape is derived from the numeric arguments.");
+  }
+
+  @Override
+  protected int getValueNumberForDTypeArgument() {
+    // 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.
+  }
+}