Unmodeled `Reshape` and `Permute` layer classes drop the output shape for the rest of the chain

[khatchad]

An unmodeled Keras layer class in the middle of a chain drops the output shape for everything downstream of it. Two such classes are reachable in ordinary models: MaxPool2D and LSTM.

Witness

The fixture tf2_test_layer_call_result_rank.py carries four chains behind four sinks, all ending in the same Dense(10) and all called with the same batch of 32.

chain	consumer types as
Dense only	(32, 10) float32
Conv2D then Flatten then Dense	(32, 10) float32
Conv2D then MaxPool2D then Flatten then Dense	{? of float32}
LSTM then Dense	{? of float32}

The first two rows are the controls and they pass, so the isolation is pinned by passing tests beside the failing ones rather than by removing code and re-running. The dtype survives in every row. What is lost is the shape, and it is lost as ⊤ rather than as a wildcard axis, so the result declares nothing about rank at all.

Inserting a single pooling layer into a chain that otherwise resolves is what turns the third row from the second.

Cause

The model file describes the layer classes Dense, Conv2D, Flatten, Embedding, Dropout and GlobalAveragePooling1D, each paired with a generator that computes the call's output shape. Neither MaxPool2D nor LSTM appears. An instance of an unmodeled class is called rather than applied, so the call produces no output shape and every value after it in the chain inherits ⊤.

This is a gap in an existing family rather than an absent capability: GlobalAveragePooling1DCall already models a rank-reducing pooling call, so the shape of the fix is established.

Why It Is Worth Prioritising

A ⊤ shape is not a coarse answer, it is the absence of one. A declaration derived from it constrains nothing, and downstream consumers that read a static axis cannot be served by it at all. Pooling after a convolution and a recurrent layer before a projection are both ordinary shapes for a model to take, so the two missing classes sit on common paths rather than exotic ones.

MaxPool2D is the more mechanical of the two: the output shape is a function of the input shape, the pool size, the strides and the padding, in the same way the existing pooling generator computes its own. LSTM needs the units argument stored on the instance and read back at the call, in the same way Dense stores units and Embedding stores output_dim, with the wrinkle that return_sequences decides whether the temporal axis survives.

[khatchad]

Widening this beyond the two classes in the title, because the two were the witnesses I happened to write rather than the extent of the gap.

The model file describes seven layer classes: Conv2D, Dense, Dropout, Embedding, Flatten, GlobalAveragePooling1D and LayerNormalization. Every other layer class is unmodeled, and each one drops the shape for the rest of its chain in the same way the two witnesses here do.

Classes that are unmodeled today and ordinary enough to sit mid-chain in a real model:

• Pooling: MaxPool2D and its aliases, AveragePooling2D, GlobalMaxPooling*
• Recurrent: LSTM, GRU, SimpleRNN, and Bidirectional as a wrapper over them
• Normalization: BatchNormalization
• Activation layers: ReLU, LeakyReLU, Softmax
• Shape-affecting: Conv2DTranspose, ZeroPadding2D, Concatenate, Reshape, Permute
• Sequence: Masking

Not all cost the same to model, and they divide by how the output shape relates to the input:

1. Shape-preserving. BatchNormalization, the activation layers (ReLU, LeakyReLU, Softmax), and Masking. The generator returns the input shape unchanged, which is the cheapest arm.
2. Computed from arguments. Pooling, Conv2DTranspose and ZeroPadding2D, where the output follows from the input shape plus pool size, strides and padding. GlobalAveragePooling1DCall is the precedent.
3. Declared width. Recurrent layers, where units is stored on the instance and read back at the call, exactly as Dense stores units and Embedding stores output_dim, with return_sequences deciding whether the temporal axis survives.

Arm 1 is worth doing first: it is the largest group by class count, the generator is trivial, and a shape-preserving layer that drops the shape is the most surprising of the three to a reader of the output, since nothing about the layer suggests it should affect the result's type at all.

[khatchad]

Status after two pull requests, so a reader does not have to diff them to learn what is left.

Landed

arm	classes	where
Shape-preserving	BatchNormalization, ReLU, LeakyReLU, Softmax, Masking, Activation	ponder-lab#758
Pooling	MaxPool2D, AveragePooling2D and their alias spellings, GlobalMaxPooling1D/2D	ponder-lab#759
Recurrent	LSTM, GRU, SimpleRNN, Bidirectional	ponder-lab#759

Both witnesses in this issue now pass as plain tests: the Conv2D then MaxPool2D then Flatten then Dense chain, and the LSTM then Dense chain.

Remaining

Conv2DTranspose, Concatenate and ZeroPadding2D are still unmodeled and still drop the shape for the rest of their chain. They are the last members reachable in ordinary convolutional code, and the mechanism is unchanged from the two witnesses above.

Reshape and Permute were listed earlier in this thread and are worth separating: both are shape-transforming rather than shape-preserving or window-computed, so each wants its own reasoning rather than joining an existing arm.

Three Findings Worth Keeping

A constructor without a __call__ is a decision, not a gap. LayerNormalization and Dropout carry constructors and deliberately no call. Modeling their calls cost a fixture six concrete shape members, because the call REPLACED whatever the value had been getting rather than adding to it. More modeling is not monotonically better.

Adding the class is not enough. Each layer also needs attribute wiring onto the layers object, a <new> plus a <putfield ... ref="layers">. Without it the class never resolves, and the symptom is a test reporting zero tensor parameters, which reads like a shape defect and is not one.

Modeling a layer can leak into the model it is called through. Modeling an activation let the shape flow along a chain, which was the point, and the same propagation typed the model parameter the layer was called through, which was not. That is #842, fixed by pinning values whose points-to set holds only Keras models. The symptom to watch is a wrong PARAMETER COUNT rather than a wrong shape.