Skip to content

The dtype value reader's fallback chain has less reach than the shape reader's, so one value resolves a shape and floors a dtype #936

Description

@khatchad

On one value, in one generator, the shape read resolves and the dtype read gives up. The elementwise generator for a residual add inside a transformer block reads its operand, the block's x parameter, twice: the shape reader walks to the callers and returns (2, 2). The dtype reader walks the same callers, dies at a tf.cast result in the calling method, floors to ⊥, and the generator seeds the result with an unknown dtype. The two halves of that chain are filed already: the floor is #620 non-aborting floor, and the empty-to-unknown step is the idiom #913 describes. Neither says that the two value readers have different reach on the same value. That is this issue, and it is the whole signature of a shape kept with a dtype lost.

Demonstrated

The repository's vendored gpt-2 fixture (gpt2_proj, transformer layers stubbed as pass-throughs, so the first-iteration dtype is int32), analysed under fine logging. In DecoderLayer.call, for the first residual add x + out (value 47), the two reads of the operand x (value 3) sit 126 lines apart in one log:

[ElementWiseOperation] EWO record broadcast for source v47 in script gpt2_model.py.DecoderLayer.call.do()LRoot; [ctx#347d4dea]: xVn 3 shapes [[D:Constant,2, D:Constant,2]], yVn 27 shapes [[D:Constant,2, D:Constant,2]].
[TensorGenerator] Could not trace dtype for value number 3 in script gpt2_model.py.DecoderLayer.call.do()LRoot; [ctx#347d4dea]; flooring to ⊥ (not a tensor). wala/ML#620.
[ElementWiseOperation] ElementWiseOperation getDefaultDTypes dtypes: []

The seed that results: [{[D:Constant,2, D:Constant,2] of unknown}] for v47, and the same for the second residual add and its intermediate (values 51 and 54). The dtype walk went from the parameter through the layer's trampoline into Gpt2.call, reached the tf.cast(x, tf.int32) call's result there, and floored. The shape walk over the same parameter resolved. Both readers start from the same points-to set and the same parameter position. The dtype reader's fallback chain (getDTypes at the parameter arm, TensorGenerator around line 6350) is shorter than the shape reader's (getShapeResultOfValue's parameter arm around line 2965, which also carries the declared-contract seed of #717 and the exact-mode remainder of #718). Which step the shape chain takes that the dtype chain lacks, on this value, is the trace this issue asks for. The log fixes the divergence to the parameter walk and the cast result.

Two things this fixture also shows, which bound what it can witness:

  • The degraded seed is repaired before it can be observed at a parameter. The engine installs a dtype-fill feed on an all-unknown-dtype seed (feed DTYPE_FILL/BROADCAST: v47), which fills the dtype from the operands' dataflow state, and the final state of all three residual values is (2, 2) int32. So DecoderLayer.call's parameter reads one member, (2, 2) int32, and the fixture cannot show a loop-carried arm with a lost dtype beside a first-iteration arm with a kept one.
  • The resolver's census for the analysis reads 233 queries and 0 in cycles. No query cycle exists here, so nothing about the resolver's cycle policy on the dtype axis applies to this loss.

A minimal fixture separately shows that a layer-call result fed back through a loop does reach the callee's parameter as a second member (hidden, present = block(hidden) over two blocks whose call casts to int64: the parameter reads {(2, 2) int32, (2, 2) int64}). The loop feedback and the tuple unpack inside it are not where a dtype is lost.

Inferred, Not Run

Programs whose transformer block is not stubbed show, in their emitted types, parameters with two members of exactly the same shape, one with a known floating dtype and one unknown. The mechanism above predicts that pair if the dtype-fill repair does not fire, or does not reach the returned value, in a real block: the unknown arm is fed back and joins the first arm at the parameter. What makes it survive there is the parameter destination's compute-dtype cast (ComputeDTypeCastOp, which rewrites a floating-point cell to the layer's compute dtype and leaves any other cell alone): a known loop-carried floating dtype is normalised onto the first arm and disappears as a distinct member, while UNKNOWN is declared non-floating and passes untouched. So the only loop-carried dtype that can appear beside the first arm is an unknown one, which is the pair observed. This is inference from the code and the minimal fixture. Those programs are not in this repository. The probe that would test it is the vendored gpt-2 fixture with the layers un-stubbed, read at the same parameter.

A Fixture That Cannot Show What It Was Built To Show

The minimal fixture's first version cast to float64. The parameter then read {(2, 2) float32} only, which reads as "the loop-carried arm never arrives" and is not that: the compute-dtype cast at the parameter folded the float64 arm onto float32, and the two arms coincided. The result is indistinguishable from a severed edge, for a reason invisible in the result. An integral cast is the version that works, because the operator leaves integral cells alone. A distinguisher the destination's own operator normalises measures nothing. Whoever builds the next loop-feedback fixture should pick one the parameter operators pass through.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions