`set_shape` XML modeling claims a return value, but `tf.Tensor.set_shape` returns `None`

[khatchad]

Observation

Per the tf.Tensor.set_shape API docs, the method returns None. The TF source confirms this — set_shape is a side-effecting mutation that updates the receiver Tensor's static shape, not a return-value producer.

tensorflow.xml claims otherwise:

<class name="set_shape" allocatable="true">
  <method name="do" descriptor="()LRoot;" numArgs="2" paramNames="self shape">
    <return value="self" />
  </method>
</class>

(tensorflow.xml:1660-1664)

This <return value="self"/> claims the call returns self. In the callable-as-attribute context, self is the callable instance — not the receiver Tensor and not None. Misleading regardless of how it's read.

Surfaced By

Reviewing the explorations in ponder-lab/ML#338 (closed): a follow-up to #550 used a similar <return value="self"/> pattern for fresh <method name="set_shape"> declarations on Tensor classes. Comparing both modelings against the real TF API surfaced the mismatch.

Possible Fix

Two reasonable options:

1. Empty body (no <return>): the WALA Python frontend can synthesize a None return for methods without an explicit <return> (verify this against existing void-returning models).
2. Explicit None allocation: <new def="r" class="Lnone"/><return value="r"/> or similar, following whatever convention tensorflow.xml uses elsewhere for None returns.

Whichever pattern lands should be applied to:

• <class name="set_shape"> at tensorflow.xml:1660-1664.
• Any new modeling added by Migrate set_shape recognition from Java IR-scan to <method> on Tensor classes in tensorflow.xml #550 implementation work.

Out of Scope

This is independent of the dispatch-mechanism debate in #550. The return-value mismatch is a separate API-correctness bug that applies regardless of whether set_shape is recognized via Java IR-syntactic scan, XML class-method dispatch, or XML callable-as-attribute.

Related

• Migrate set_shape recognition from Java IR-scan to <method> on Tensor classes in tensorflow.xml #550 — set_shape recognition architecture; would consume any cleanup here.
• ponder-lab/ML#333 — Java-side recognition; doesn't touch the XML body but uses getSetShapeCallsSyntactic to find the calls.
• ponder-lab/ML#338 (closed) — surfaced the mismatch during XML-side exploration.

[khatchad]

Quick spike on the suggested fix paths:

• Empty body (option 1): removing the <return value="self"/> from the legacy <class name="set_shape">.do() at tensorflow.xml:1652-1656 produces 761 test errors. WALA's call-graph dispatcher or XML loader appears to require at least one body element on <method> declarations; the parser tolerates the empty body but the dispatch downstream doesn't.
• Explicit None allocation (option 2): no precedent for None-allocation in tensorflow.xml's current modeling. Would require introducing a sentinel class like <class name="None" allocatable="true"> and routing all void-returning APIs through it. That's a broader convention change, not a localized fix.

Leaving the current <return value="self"/> in place. It's API-wrong but observably inert: the legacy getShapeSourceCalls(set_shape, ...) dispatch was removed by ponder-lab/ML#333 in favor of getSetShapeCallsSyntactic, so no consumer reads the def's PTS for these calls. The right fix is the broader void-return convention, not a one-site patch.

[khatchad]

Started investigating this and the dispatch model is more entangled than the issue's proposed-fix options suggest. Posting findings before picking a path.

Why the current <return value="self"/> Exists

The set_shape class at tensorflow.xml:1652-1656 is the bound-method object allocated and stored on Tensor instances via the <putfield> chain around line 2128 (<new def="y" class="Ltensorflow/functions/set_shape"/><putfield ... field="set_shape" ... value="y"/>). The do method on this class is the call-handler that fires when a user writes x.set_shape(shape). Returning self returns the bound-method object — not the receiver x and not None. In the callable-as-attribute context this is the standard pattern for "pass through the callable" but it's semantically wrong here because set_shape is side-effecting and returns None in real TF.

Why the Issue's Options Don't Cleanly Apply

Option 1 (Empty body, no <return>): untested against this loader path. The __init__ precedent at tensorflow.xml:3250 uses descriptor="()V" AND has no <return> — those go together. Just deleting <return> while keeping descriptor="()LRoot;" produces an inconsistent declaration ("returns LRoot" but no <return> emitter). Behavior is loader-defined.

Option 2 (Explicit None allocation): the proposed class="Lnone" type doesn't exist in PythonTypes.java (no kNone / globalNone either). I checked: Lnone doesn't appear anywhere in the cast.python source or in any *.xml summary. So adopting it would require either inventing the type or substituting Lobject (an existing dummy type), which mis-declares the call's return type as Object and would silently flow into downstream tensor analyses that test instanceof PythonTypes.None (none exist today, but the misdirection compounds).

Possible Third Option

Switch the class itself: declare set_shape's do with descriptor="()V" and drop <return>, matching the __init__ precedent exactly. The risk is that do is special-cased by the loader (callable-as-attribute dispatch) and may not tolerate a V-typed declaration. The set_shape block is registered to the Ltensorflow/functions/... package, which is the standard pattern for callable-as-attribute classes — every other such class in tensorflow.xml uses ()LRoot;, so we'd be the first.

Open Empirical Questions

1. Whether the Python frontend distinguishes "method returned LRoot via <return value=X/>" from "method returned LRoot via fallthrough" at all. If not, all three options produce identical IR.
2. Whether any consumer of x.set_shape(s)'s return exists (i.e. any place where the result is bound: y = x.set_shape(s)). If none, all three options are equivalent in practice and the fix is cosmetic. A grep for \b\w+\s*=\s*\w+\.set_shape\b against the test corpus would settle it.

[khatchad]

Answering question 2 empirically:

$ grep -rnE '\.set_shape\(' com.ibm.wala.cast.python.test/data/
com.ibm.wala.cast.python.test/data/train.py:53:    imagei.set_shape([IMG_DIM * IMG_DIM * IMG_DIM])
com.ibm.wala.cast.python.test/data/train.py:56:    labels.set_shape([IMG_DIM * IMG_DIM * IMG_DIM])
com.ibm.wala.cast.python.test/data/tf2_test_set_shape_non_tensor.py:20:fake.set_shape([1, 2, 3])

$ grep -rnE '\b\w+\s*=\s*\w+\.set_shape\(' --include='*.py' .
(no matches)

All three call sites in the test corpus are statement-position invocations, not return-binding. Zero consumers of set_shape's return value across the repo today.

So the fix here is cosmetic for our current corpus — any of the three options I outlined produces the same observable analysis result. The benefit is forward-looking: if a future test fixture writes y = x.set_shape(s) expecting y to be None, the current <return value="self"/> would silently flow the set_shape callable through y, masking the bug.

That tips toward landing the descriptor="()V" + empty-body fix (option 3) as a small, principled correction. Will sit on it for a day in case @msridhar / @liblit weighs in on whether do is special-cased in the loader.