Resolve tf.TensorShape literals in shape arguments

com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/tensorflow/v2/TestDatasets.java

@@ -919,13 +919,10 @@
 
   /**
    * The {@code tf.TensorShape} literal form of {@link #testDatasetFromGeneratorDeclaredShapes()}:
-   * the declared shape is wrapped in a {@code TensorShape} constructor rather than a plain tuple.
-   * {@code tf.TensorShape} is unmodeled, so the argument is present but unparseable and the shape
-   * axis soundly degrades to ⊤ ({@code {? of int32}}) instead of composing {@code (2,)}.
-   *
-   * <p>TODO: Flip to expect {@code TensorType.of(INT_32, 2)} when <a
-   * href="https://github.com/wala/ML/issues/789">wala/ML#789</a> parses {@code TensorShape}
-   * constructor literals in shape arguments.
+   * the declared shape is wrapped in a {@code TensorShape} constructor rather than a plain tuple,
+   * whose stored {@code dims} the shape-argument parser recurses into (<a
+   * href="https://github.com/wala/ML/issues/789">wala/ML#789</a>), composing {@code (2,)} exactly
+   * like the tuple forms.
    *
    * @throws ClassHierarchyException On WALA class-hierarchy error.
    * @throws IllegalArgumentException On illegal argument.
@@ -940,7 +937,7 @@
         "consume",
         1,
         1,
-        Map.of(2, Set.of(new TensorType(INT_32, null))));
+        Map.of(2, Set.of(TensorType.of(INT_32, 2))));
   }
 
   /**

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

@@ -197,6 +197,8 @@
         <putfield class="LRoot" field="random" fieldType="LRoot" ref="Dataset" value="dsrandom"/>
         <new def="from_tensors" class="Ltensorflow/data/Dataset/from_tensors"/>
         <putfield class="LRoot" field="from_tensors" fieldType="LRoot" ref="Dataset" value="from_tensors"/>
+        <new def="TensorShape" class="Ltensorflow/framework/TensorShape"/>
+        <putfield class="LRoot" field="TensorShape" fieldType="LRoot" ref="x" value="TensorShape"/>
         <new def="TensorSpec" class="Ltensorflow/framework/TensorSpec"/>
         <putfield class="LRoot" field="TensorSpec" fieldType="LRoot" ref="x" value="TensorSpec"/>
         <new def="RaggedTensorSpec" class="Ltensorflow/framework/RaggedTensorSpec"/>
@@ -1105,6 +1107,14 @@
       </class>
     </package>
     <package name="tensorflow/framework">
+      <class name="TensorShape" allocatable="true">
+        <!-- https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/TensorShape. The constructor stores its dims argument on the result so the shape-argument parser can recurse into the wrapped list (wala/ML#789), mirroring TensorSpec's shape field. -->
+        <method name="do" descriptor="()LRoot;" numArgs="3" paramNames="self dims name">
+          <new def="ts" class="Ltensorflow/framework/TensorShape"/>
+          <putfield class="LRoot" field="dims" fieldType="LRoot" ref="ts" value="dims"/>
+          <return value="ts"/>
+        </method>
+      </class>
       <class name="TensorSpec" allocatable="true">
         <method name="do" descriptor="()LRoot;" numArgs="4" paramNames="self shape dtype name">
           <new def="spec" class="Ltensorflow/framework/TensorSpec"/>

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

@@ -619,7 +619,8 @@
               if (innerReference.equals(tuple)
                   || innerReference.equals(list)
                   || innerReference.equals(TensorFlowTypes.TENSOR_SPEC)
-                  || innerReference.equals(TensorFlowTypes.RAGGED_TENSOR_SPEC)) {
+                  || innerReference.equals(TensorFlowTypes.RAGGED_TENSOR_SPEC)
+                  || innerReference.equals(TensorFlowTypes.TENSOR_SHAPE)) {
                 // Nested tuple/list or Spec. Recurse.
                 Set<List<Dimension<?>>> nestedShapes =
                     this.getShapesFromShapeArgument(
@@ -782,17 +783,21 @@
         if (constantShapes == null) return null;
         ret.addAll(constantShapes);
       } else if (reference.equals(TensorFlowTypes.TENSOR_SPEC)
-          || reference.equals(TensorFlowTypes.RAGGED_TENSOR_SPEC)) {
-        // We have a TensorSpec or RaggedTensorSpec. These objects carry shape and dtype
-        // information in their fields. We extract the 'shape' field and recurse to
-        // parse the actual shape structure (usually a tuple or list of integers).
+          || reference.equals(TensorFlowTypes.RAGGED_TENSOR_SPEC)
+          || reference.equals(TensorFlowTypes.TENSOR_SHAPE)) {
+        // We have a TensorSpec, RaggedTensorSpec, or TensorShape. These objects carry their shape
+        // structure in a field ('shape' for the specs, the stored 'dims' argument for a
+        // TensorShape constructor, wala/ML#789); extract it and recurse to parse the actual
+        // structure (usually a tuple or list of integers).
         IField shapeField =
             builder
                 .getClassHierarchy()
                 .resolveField(
                     reference.equals(TensorFlowTypes.TENSOR_SPEC)
                         ? TensorFlowTypes.SPEC_SHAPE
-                        : TensorFlowTypes.RAGGED_SPEC_SHAPE);
+                        : reference.equals(TensorFlowTypes.RAGGED_TENSOR_SPEC)
+                            ? TensorFlowTypes.RAGGED_SPEC_SHAPE
+                            : TensorFlowTypes.TENSOR_SHAPE_DIMS);
         PointerKey shapePK = builder.getPointerKeyForInstanceField(instanceKey, shapeField);
         OrdinalSet<InstanceKey> shapePts = pointerAnalysis.getPointsToSet(shapePK);
         if (shapePts == null || shapePts.isEmpty()) return null;

com.ibm.wala.cast.python.ml/source/com/ibm/wala/cast/python/ml/types/TensorFlowTypes.java

@@ -329,6 +329,15 @@ public boolean canConvertTo(DType other) {
 
   public static final String RAGGED_TENSOR_SPEC_SIGNATURE = "tf.RaggedTensorSpec()";
 
+  /** https://www.tensorflow.org/api_docs/python/tf/TensorShape. */
+  public static final TypeReference TENSOR_SHAPE =
+      TypeReference.findOrCreate(
+          pythonLoader, TypeName.findOrCreate("Ltensorflow/framework/TensorShape"));
+
+  /** The {@code dims} argument the {@code TensorShape} constructor stores (wala/ML#789). */
+  public static final FieldReference TENSOR_SHAPE_DIMS =
+      FieldReference.findOrCreate(TENSOR_SHAPE, findOrCreateAsciiAtom("dims"), Root);
+
   public static final FieldReference SPEC_SHAPE =
       FieldReference.findOrCreate(TENSOR_SPEC, findOrCreateAsciiAtom("shape"), Root);