Tensor-type analysis misses parameters that receive tensors at their call sites

[khatchad]

Summary

The tensor-type analysis does not associate a TensorType with function parameters that are bound to tensor values at their call sites. Affected functions are consequently treated as having no tensor parameter by downstream consumers.

Verified Cases (Call-Site Evidence)

Reading each call site confirms the parameter receives a tensor, yet the analysis yields no tensor type:

Function	Param	Bound to, at the call site
Gpt2.get_loss(self, real, pred)	real, pred	get_loss(targets, predictions): model output + label tensors
BiLSTM.call(self, inputs)	inputs	token-ID tensor into a Keras Embedding
TextCNN.predict(self, inputs)	inputs	token-ID tensor into a Keras Embedding
TuckERLoader.target_convert(self, targets, ...)	targets	a tf.data padded_batch int32 element

Three propagation shapes: interprocedural (a tensor argument to a callee), Keras call dispatch (a tensor into __call__ reaching an Embedding), and a tf.data dataset element.

Suggested Minimal Repro (Interprocedural)

import tensorflow as tf

class Model:
    def call(self, x):
        return x * 2.0

    def get_loss(self, real, pred):
        return tf.reduce_mean(tf.square(pred - real))

    def train_step(self, inputs, targets):
        predictions = self.call(inputs)
        return self.get_loss(targets, predictions)  # both args are tensors

def main():
    Model().train_step(tf.constant([1., 2., 3.]), tf.constant([1., 1., 1.]))

Expected: get_loss's real and pred each carry a TensorType(float32). Observed: neither parameter receives a tensor type.

Context

Surfaced while regenerating a @tf.function refactoring corpus, where these functions are left undecorated because the parameter is not seen as a tensor. An earlier evaluation classified some of these as tensor parameters, but that pipeline predates a rewrite of the consumer's classification code, so this is reported as a current under-approximation, not a verified version-to-version regression.

[khatchad]

A related shape: parameters that are a container of tensors (e.g. a model called as model([a, b, c]), where the parameter binds to a list of tensors).

The layer is undetermined here as well. The consumer's container detection (Phase 3) recognizes list/tuple/dict/set/enumerate and tests whether the parameter's points-to set includes a container instance carrying an element tensor. That spans both the analysis (the container's tensor membership and the points-to relation) and consumer-side code (resolving the parameter to a pointer key and issuing the points-to query, the same class of binding whose mis-resolution can drop a type in the direct-tensor cases). So a miss can originate in either layer and isn't determinable without running the analysis.

Recording it as a second shape of the same symptom, not a separately attributed bug.

[khatchad]

Update: the suggested minimal interprocedural repro types both parameters correctly in the current analysis. With get_loss made reachable (a module-level Model().train_step(...)) and querying Model.get_loss, both real and pred resolve to (3,) float32.

Minimal repros of the other two propagation shapes also type the parameter precisely:

• Keras call dispatch: a token-id tensor passed to call(inputs) that feeds an Embedding resolves inputs to (1, 3) int32.
• tf.data element: a dataset element passed to a function resolves the parameter to (2,) int32.

So in isolation these three shapes are handled, and the originally-reported miss is not explained by a generic failure of interprocedural, Keras-call, or tf.data-element propagation. Localizing it needs a self-contained reproducing snippet (the parameter seen as non-tensor) and a check of whether the gap is in the analysis or the downstream classifier; the likely cause is model-specific structure (lazy build() as in #595, cross-module resolution, decorators).

The three working shapes are now guarded by regression tests in ponder-lab#453.

[khatchad]

Update from instrumented analysis on a real multi-module subject. The minimal repros all type the parameter correctly, so this isolates the trigger.

For a failing case, a loss method's parameter that receives a model output and a dataset-threaded label:

Reachability is fine. The owning function is in the call graph, and the entire tensor-producing driver is analyzed end to end: module top-level → training function → a custom fit → train_step → the loss method. Not an entry-point or reachability gap.

Emission is the gap. No per-parameter TensorType is associated with the parameter's pointer key (verified by node identity, independent of any downstream source-position matching). And it's the first hop: train_step's own inputs/targets parameters are themselves untyped, so tensor types aren't crossing from the tf.data elements into the callee parameters; the absence then cascades to the loss method.

Triggering shape (not exercised by the minimals): a custom fit (not keras.Model.fit) iterates a tf.data dataset (here experimental_distribute_dataset-wrapped) and passes the yielded (inputs, targets) tuple into train_step(self, inputs, targets), which computes predictions = self(inputs) and calls loss(targets, predictions). The direct-call minimal types the parameter; the iterate-in-fit-then-pass-the-tuple multi-hop does not.

This is a current under-approximation, not established as a version regression. Suggested next step: a repro of exactly that shape, checking whether train_step's parameters receive a per-parameter TensorType, to pin whether the type drops at the dataset-element typing or at the call into the callee.