`tf.gather` returns ⊤ when the indices' shape is unknown, discarding the table's known trailing axes

[khatchad]

Compute tf.gather's output shape instead of passing the table through (#727, fixing #815) made tf.gather return ⊤ whenever the indices' shape is unknown, discarding the trailing axes of the table, which are known. Shapes that resolved at 0.52.77 are unknown-rank at 0.52.78.

The Code Path

Gather extends EmbeddingLookup, whose getDefaultShapes derives ids.shape + params.shape[1:]:

if (paramsShapes == null || idsShapes == null) return null;

An unknown ids therefore poisons the whole result, including params.shape[1:], which does not depend on ids at all.

The Pattern

table  = tf.zeros((100, 16))
idx    = something_the_analysis_cannot_resolve()   # rank unknown
picked = tf.gather(table, idx)                     # truth: idx.shape + (16,)

Before #727 the model passed the table through, so picked reported the table's own shape, rank 2 with a known trailing axis. After it, picked is ⊤.

What Makes This Awkward

The pass-through was not right in general: it named the table's leading extent where the truth is the indices' extent. It happened to read correctly whenever the table's leading axis was itself unresolved, since both then print the same way. So the previous answer was accidentally right and the new one is honestly unknown, which is not the usual regression shape.

That rules out the obvious repair. Assuming ids is rank 1, which it usually is, would under-rank every call that passes a higher-rank index tensor, and an argument of rank 3 is no more a subtype of a rank-2 specification than a rank-1 one is. The reasoning #813 used against over-ranking applies in both directions.

Two Directions That Are Sound

Resolve the indices. The common idiom builds them by slicing a column out of an edge list, idx = pairs[:, 0], which is exactly what #813's machinery types as rank 1. When that resolves, gather yields the correct shape on its own terms rather than by accident, which is better than what the pass-through gave.

Stop the whole-set poisoning. EmbeddingLookup is on the legacy null-means-⊤ path, and its record counterpart is ShapeResult.fromLegacy(getDefaultShapes(...)), so it inherits the all-or-nothing behavior that ShapeResult (#718) exists to replace. If ids resolves to some members and not others, crossing the resolvable ones and carrying the rest as the unknown remainder would preserve what is known. Whether ids is partially resolvable in the reported case is not established; that needs a FINE probe, and if ids is wholly ⊤ there is no partial to preserve and this direction buys nothing.

The first direction is the one that addresses the cause. The second is worth doing regardless, since the poisoning is a contract violation independent of this symptom.

Bisected

The behavior is absent at 0.52.77 and present at 0.52.78, 0.52.79 and 0.52.80, with every other input held fixed. It affects tf.nn.embedding_lookup on the same code path.

[khatchad]

Probed. The indices are not partially resolvable, so the partial-results direction in the description buys nothing here.

What The Probe Shows

A fixture mirroring the reporting idiom, with the assertions holding at run time:

Value	Truth	Inferred
adjacency[:, 0], receiver a parameter	int32 (30,)	int32 (30, 2)
tf.gather(table, that)	float32 (30, 16)	float32 (30, 2, 16)
adjacency_list[:, 0], through enumerate	int32 (30,)	no tensor type at all
tf.gather(table, that)	float32 (30, 16)	⊤

The reported case is the third row: the index is wholly untyped rather than partly resolved, so there are no members to cross and no remainder to carry. Rows one and two are the same defect one form earlier, where the index resolves but at the receiver's rank, and the gather composes that faithfully into a rank-3 result.

Consequence For This Issue

The gather model is correct in every row. It composes exactly what it is handed, and what it is handed is wrong or absent.

The whole-set poisoning noted in the description is still a contract violation worth fixing on its own terms, but it is not what produces the reported symptom.