Token-mask helpers infer unknown rank where the body requires rank 2

[khatchad]

Two gpt-2 helpers are inferred with unknown rank while their bodies determine the rank exactly.

• get_padding_mask(seq) slices seq[:, tf.newaxis, tf.newaxis, :], which is well-formed only for a rank-2 argument.
• create_masks(inp) indexes tf.shape(inp)[1] and forwards inp to get_padding_mask, so it inherits the same requirement.

Both are inferred as unknown-rank rather than as a rank-2 tensor with two variable axes. The axes themselves are correctly variable, since the sole call site passes the pipeline token tensor whose batch and sequence lengths both vary, so the recovery wanted here is the rank, not the extents.

Rank is derivable from the body in both cases: an extended slice with two tf.newaxis insertions and two positional slices constrains the operand rank directly, and a tf.shape(x)[1] read establishes a lower bound. Recovering it would move these parameters from an unknown-rank spec to a rank-2 spec with two required wildcards, which is a strictly better signature.

The call sites and line numbers behind each axis are recorded alongside the analysis notes.

[khatchad]

These Two Are Precision, Not Correctness

Applied the static-versus-dynamic test here as well, since an unresolved rank matters more when something consumes it.

get_padding_mask  emitted shape=None      -> traces OK
                  rank-2 (None, None)     -> traces OK

The body is seq[:, tf.newaxis, tf.newaxis, :], and TensorFlow accepts that under an unknown-rank spec: the slicing is resolved at graph-construction time against the actual tensor rather than requiring a statically known rank. create_masks forwards to it and reads tf.shape(inp)[1], which is dynamic and equally safe.

So recovering the rank here is worth doing for tightness, and nothing depends on it for correctness. That distinguishes this issue from #809 and #811, where one axis in each turned out to be load-bearing: in those the layer reads the dimension off the static shape and hands it to a weight shape, which needs a Python integer.

Recording it because the three issues were filed together from the same corpus sweep and it would be easy to assume they carry the same urgency. This one does not.