fix(video): declared-shape predict for 20 video models and videomae pretraining fidelity

[ooples]

Summary

The follow-up work PR #2174 listed as "what this does not do". Four commits, all root-cause fixes in the library, each with a regression test that was proven to fail against the unfixed source (method below).

This PR is based on fix/videomae-pretraining-and-timesformer (PR #2174), which is still open at 4dedc7599. It is not based on master; merge #2174 first, or retarget this once #2174 lands. Every line number below is in the unfixed file at that base commit.

Four groups of defects:

1. 13 optical-flow models could not Predict on their own declared input shape — their shared base rejected every unbatched input.
2. Seven more video models failed the same declared-shape probe for model-specific reasons; five are fixed, four are a decision for you (own section below).
3. TimeSformer ignored Architecture.InputFrames, and had two output-convention bugs (one of them a real numerical bug: a batch softmax).
4. VideoMAE's masked-autoencoder pretraining never trained, leaked masked content into the encoder's gradient, severed its own decoder tape, and had no per-patch target normalisation.

---

1. Optical flow: Predict on the declared shape

Root cause: src/Video/OpticalFlowBase.cs:331 threw "Input must be rank 4 [batch, 2*channels, height, width], got rank 3." for any rank-3 input. Thirteen models declare their architecture as InputType.ThreeDimensional with InputDepth = 6 (a frame PAIR stacked on the channel axis), so GetInputShape() is the unbatched [6, H, W] — the shape the model advertises is exactly the shape its Predict refused.

Affected: DKM, DPFlow, FlowDiffuser, FlowFormer++, MemFlow, NeuFlowV2, RoMa, RPKNet, SEA-RAFT, SKFlow, UFM, UniMatch, VideoFlow. RAPIDFlow hit the same check with an unbatched pair.

Fix (PredictCore): a rank-3 [2*C, H, W] pair is promoted with the established shared helper NeuralNetworkBase.PromoteToBatchedTensor, and the unit batch axis is squeezed back off the flow with a recorded Engine.Reshape — unbatched in, unbatched out, the same rule the base PredictCore applies to every model that inherits it (NeuralNetworkBase.cs:6385 promote / :6434 squeeze). The rank-4 body moved into a private PredictBatchedPair, so the batched path is byte-for-byte the same computation.

The declaration was brought in line with the code: both family [TensorLayout]s are now BatchOptional = true, and OutputAxesFor answers rank 3 ([Channels=Fixed(2), Height=Same, Width=Same]). Without that, the base would have accepted a rank the type does not declare.

RAFT (src/Video/Motion/RAFT.cs:267) overrides PredictCore but inherits that layout, so it owes the same contract: it sliced input.Shape[3] unconditionally and threw IndexOutOfRangeException on an unbatched pair. It now promotes/squeezes the same way, and a non-3/4 rank gets a clear ArgumentException instead of an index crash.

	before	after
13 pair-declaring models, Predict(GetInputShape())	ArgumentException: Input must be rank 4 …	[2, H, W], equal element-for-element to the batched [1, 2, H, W]
RAPIDFlow, unbatched pair [6, H, W]	same ArgumentException	[2, H, W]
RAFT, unbatched pair [6, H, W]	IndexOutOfRangeException	[2, H, W]

---

2. The seven model-specific declared-shape failures

Model	Verdict	What changed
BasicVSR++	Predict wrong — src/Video/Enhancement/BasicVSRPlusPlus.cs:542 read [C,H,W] as a clip of C two-dimensional frames, so the feature extractor got a rank-2 "frame" (ConvolutionalLayer expects rank-3 … got rank 2)	EnhanceVideo treats [C,H,W] as the one-frame clip the super-resolution family layout already documents ("A single low-resolution frame - the degenerate one-frame clip"), runs it as [1,C,H,W] and returns [C, sH, sW]. Recorded reshapes, so training through a single frame keeps its gradient path
OpenSora	Predict wrong — src/Video/Generation/OpenSora.cs:350; the denoiser indexes axis 3 of its features, so a rank-3 latent threw IndexOutOfRangeException	denoises a rank-3 latent as a batch of one and returns it without the batch axis, the same promotion the model already applies to single images elsewhere
ProPainter	Predict wrong — src/Video/Inpainting/ProPainter.cs:270; the image path indexes axis 3, IndexOutOfRangeException	runs a rank-3 frame as a one-frame clip (frames are reconstructed independently — the frame axis acts as the batch) and squeezes back
VideoCLIP	declared shape wrong — src/Video/Understanding/VideoCLIP.cs:187 declared one [3,224,224] frame although the model's input layout is [Frames, C, H, W] and it owns a clip length (numFrames, default 32); EncodeVideo → AddBatchDimension5D threw IndexOutOfRangeException	default ctor declares FourDimensional with inputFrames: 32; numFrames (:238) is now resolved from the architecture when it declares one, with a conflicting explicit value rejected — the same rule VideoMAE and TimeSformer use
RAFT	unresolved (see next section)	only the rank-3 pair promotion from item 1
GMFlow	unresolved (see next section)	nothing
RIFE	unresolved (see next section)	nothing

---

⚠ DECISION NEEDED: four models whose declared shape cannot be fixed without a contract change

RAFT, GMFlow, RAPIDFlow and RIFE still fail the declared-shape probe, deliberately. All four document Architecture.InputDepth as the per-frame channel count, so their declared shape is ONE frame ([3, H, W]) while Predict needs a stacked PAIR ([6, H, W]). GetInputShape() is derived from the architecture, so there is no way to make the declared shape the pair without redefining what InputDepth means for these models — a public constructor argument's meaning.

The evidence that this is a deliberate, documented contract and not an oversight:

• RAFT — _channels = architecture.InputDepth (RAFT.cs:176); its <example> passes inputDepth: 3; VideoExtendedIntegrationTests.RAFT_InputChannels_DefaultsToThree asserts InputChannels == 3 for a depth-3 architecture; the generated fixture builds it with inputDepth: 3 and feeds [1,6,64,64]; SEARAFT.cs:78-81 explicitly excludes RAFT and GMFlow from the "declare depth 6" rule because they have a separate 3-channel context encoder.
• GMFlow — same split (GMFlow.cs:153), same inputDepth: 3 example and fixture.
• RAPIDFlow — TryGetArchitectureInputShape() => null with a doc comment stating "Architecture.InputDepth itself reports the SINGLE-FRAME count (3)", and TestScaffoldGenerator.cs:9545-9547 states the same in the fixture comment.
• RIFE — this is the frame-interpolation family's documented rule, not RIFE's own: FrameInterpolationBase.cs:171-184, RIFE's ctor comment, its <example>, and the generator's fixture note all define InputDepth as per-frame. Changing it touches ~21 models (and FILM's fixture already contradicts the family doc by passing inputDepth: 6, so the family is already inconsistent).

Changing InputDepth to mean the stacked pair for these four is a coherent option (13 of the 17 flow models already do it, and OpticalFlowBase's note now says "most of the family"), but it would flip a documented public contract, break the assertions above, and for RIFE reverse a whole family's rule. Per the "user has final say on scope decisions" rule I stopped at evidence + recommendation. Recommended: adopt "InputDepth = 2×channels" for RAFT, GMFlow and RAPIDFlow (changing _channels = InputDepth / 2 and the fixtures in lockstep so parameter counts do not move), and treat the frame-interpolation family separately as its own decision.

---

VideoCLIP: a latent bug the declared-shape change exposed

With any FourDimensional architecture — not just the new default — the base lazy-shape walk (NeuralNetworkBase.ResolveLazyLayerShapes, :5542) starts from a rank-5 clip that VideoCLIP's first convolution rejects, so the stack is left unresolved until the first forward and ParameterCount under-reports at construction:

default new VideoCLIP<float>()	ParameterCount at construction
before (declared one frame, walk happened to work)	211,126,273
after the declared-shape change, without the walk fix	38,004,480
after the walk fix (this PR)	211,126,273

That matters beyond cosmetics: ParameterCount is what the weight-streaming auto-detect, ShouldUseStreamingTraining and ConfigureInferenceForScale gate on — the base class's own comment on that walk calls out that an under-reported count "blinded every ParameterCount-gated memory decision". VideoCLIP now overrides TryGetArchitectureInputShape() to return one frame [1, C, H, W], because its layer stack genuinely runs per frame (the spatial encoder consumes frames, not clips). Precedent for the override: RAPIDFlow, UFM, DPFlow, MiDaS, FrameInterpolationBase, Transformer, PANNsModel.

Regression test VideoCLIP_DeclaringTheClip_DoesNotChangeTheResolvedLayerStack (small config): before 2,304 parameters vs 50,477 for the equivalent one-frame declaration; after, both 50,477.

---

3. TimeSformer

Frame count. src/Video/ActionRecognition/TimeSformer.cs:223 set _numFrames = numFrames (default 8) and ignored Architecture.InputFrames, so a 16-frame architecture built an 8-frame model: NumFrames and metadata said 8, the positional table was sized 8 * patches + 1, and — because the blocks are constructed with that count — TimeSformerBlockLayer.Forward(input) (the frame-count-less overload used whenever the layer is run on its own) grouped a 16-frame token sequence into 8 frames. It now applies VideoMAE's rule; the ONNX constructor's hard-coded _numFrames = 8 (:263) reads the architecture too.

The rule moved into a new internal src/Video/VideoClipFrameCount.cs shared by VideoMAE, TimeSformer and VideoCLIP: the architecture's declared count wins; an explicit non-default numFrames that disagrees throws ArgumentException(nameof(numFrames)). VideoMAE's behaviour is unchanged (it delegates to the helper; its 10 existing tests still pass).

Unbatched output. Classify returned [1, NumClasses] for a [T,C,H,W] clip. Its own XML doc says [NumClasses], the model's output [TensorLayout] is BatchOptional, and the base PredictCore squeezes the batch it adds — three statements of the same convention. Now squeezed (tokenized path only; a caller-supplied layer stack is left untouched).

Batch softmax (numerical bug). TimeSformer.cs:401's hand-rolled softmax took the max and the sum over every element of the logits tensor, so for a batch of B clips the whole [B, NumClasses] matrix summed to 1 rather than each row: each clip's "probabilities" were scaled by how confident the other clips in the batch happened to be. Measured on a 3-clip batch: row sum 0.325 instead of 1.0. Replaced with Engine.Softmax over the class axis.

Also moved the native constructor's XML docs, which were attached to the parameterless constructor (two <summary> blocks in a row), onto the constructor they describe, and documented numFrames and options.

Default parameter count unchanged: 113,541,120.

---

4. VideoMAE pretraining fidelity (Tong et al. 2022)

(a) PretrainMAE never trained. src/Video/ActionRecognition/VideoMAE.cs:345 computed a reconstruction loss and returned it — no backward, no optimizer step. The class docs, the <example>, and the name all describe pretraining, so it is meant to train. It now records the masked forward and the loss on a GradientTape and steps through BackwardAndStepOnPrecomputedLoss — the library's caller-owned-tape training entry point used by NeRF, TVAE, TabDDPM and the GANs — with the constructor's optimizer or the model's default Adam. The loss is now built from engine ops (subtract → square → multiply by a 0/1 masked-patch tensor → sum → scale) so it back-propagates; ComputeReconstructionLoss returns that same objective's value under a NoGradScope.

Measured on a fixed clip (small config: 4 frames, 32×32, 16 features, mask ratio 0.5, default Adam lr 1e-3):

	before	after
reconstruction head parameters after 60 steps	unchanged	changed
loss over 40 steps	flat (varies only with the random mask)	0.994 → 0.893, falling monotonically

Which layers update (one step, verified per layer): patch embedding (0), all 12 encoder blocks (1-12), all 4 decoder blocks (16-19) and the reconstruction head (20) — and not the classification head (13 feature-reduce, 15 classifier), which is not on the pretraining path. Before: none of them.

(b) Masked tokens were zeroed, not dropped. VideoMAE.cs:743 wrote zeros into the patch embedding in place, once. Two defects followed:

• The 3×3 encoder blocks immediately refilled every masked position from its visible neighbours plus the block bias, so from the second block on the encoder was processing masked positions as ordinary tokens (measured: masked patch feature = 1.7e-3, not 0).
• The in-place write is invisible to the autodiff tape, so the backward pass still routed gradient through the embedding's masked outputs into the patch-embedding weights — weighted by the masked patches' own pixels, i.e. the content the model is supposed to predict without seeing.

The paper's encoder drops masked tokens; a convolution over the patch grid cannot drop grid positions. So it is made sparse the way masked-image-modelling work on convolutional encoders does it (SparK, Tian et al. 2023; ConvNeXt V2's FCMAE, Woo et al. 2023): a recorded visibility mask multiplies the patch embedding and the output of every encoder block. A masked position then holds exactly zero at every block input — it contributes nothing to its visible neighbours (identical to zero padding) and never accumulates a value or a gradient of its own.

Deliberate deviation from the paper: the learned mask token is not added. It would be a new trainable parameter (and would change every default parameter count), so masked positions enter the decoder as a fixed zero "mask token" and the decoder's convolutions fill them from the visible context. This is documented on EncodeVisiblePatches and in the class remarks rather than left implicit.

(c) Tape-severing GELU in the decoder. VideoMAE.cs:792 stacked a Tensor.Transform-based GELU on every decoder block's own ReLU. Transform records no autodiff node, so it severed the tape after every block: a reconstruction loss could reach the head and nothing before it — no decoder block, no encoder block, no patch embedding. It was also numerically redundant (GELU of a non-negative ReLU output). Removed — the identical defect had already been fixed in the encoder on the base branch.

(d) No per-patch target normalisation. The target was raw pixels (ComputeReconstructionLoss, :828). New VideoMAEOptions.NormalizeTarget, defaulting to true — the paper's normlize_target default, inherited from MAE (He et al. 2022, §4) — normalises every channel of every tubelet patch over its own tubeletSize × 16 × 16 pixels as (x - mean) / (std + 1e-6) with the unbiased standard deviation, exactly the reference implementation's b (t h w) (p0 p1 p2) c view reduced over p0 p1 p2. Setting it to false restores the raw-pixel target bit-exactly (the existing exact-value test now pins that path). Documented on the option, on ComputeReconstructionLoss, on the target builder and in the class remarks.

No parameters are added anywhere: default VideoMAE parameter count unchanged at 65,788,816.

---

Test results (net10.0, Release)

Baseline = this branch's base commit 4dedc7599; "after" = this PR's head. Same machine, same filter.

Suite	Before	After
Generated model-family suites: VideoMAE, TimeSformer, VideoCLIP, all 17 optical-flow models, RIFE, BasicVSR++, ProPainter	783 passed, 23 skipped, 1 failed	783 passed, 23 skipped, 1 failed — per-class counts identical
TracedChainValidationTests	3 passed, 1 failed (VideoMAE)	3 passed, 1 failed (VideoMAE)
New regression tests (35)	33 failed, 2 passed	35 passed
VideoExtendedIntegrationTests	–	83 passed
RAPIDFlowReviewRegressionIntegrationTests	–	3 passed
ShapeContractConformanceTests	–	28 passed
TensorLayoutRankTests	–	12 passed

Neither failure is a regression. Both are the same single test, TracedChainValidationTests.TracedValidationClearsTheFalsePositivesTheLinearReadingProduces(modelName: "VideoMAE"), failing identically before and after with System.ArgumentOutOfRangeException : Inner left-operand block exceeds its span. (Parameter 'aBase'). That is the one-frame default-constructor crash owned by the never-run-tests PR; PR #2174 already documented it and deliberately kept clear of those lines, and so does this PR.

The 2 new tests that pass before the fix are guards, not evidence: TimeSformer's "no declared frame count keeps the numFrames argument" and the options default (which reads the stub described below).

ModelFamilyLawTests.DoTheMembersOfAFamilyShareOneShapeLaw timed out locally after its own 30-minute limit, so I have no before/after for it. It is the whole-model-zoo sweep (~900 constructions) that is known to time out unreliably on a local box; it was not part of the baseline run either, so I am reporting it as unmeasured rather than as a pass or a failure. CI is the arbiter for that one.

How the fails-before numbers were produced

git stash push --include-untracked -- src (source only, tests left in place), then two API-surface-only stubs re-applied to the unfixed source so the new tests compile against the old logic:

• VideoMAE.EncodeVisiblePatches private → internal (visibility, no behaviour),
• an unread VideoMAEOptions.NormalizeTarget property (the unfixed model never looks at it).

Then a full rebuild and run. Afterwards the stash was popped and the restored working tree was verified byte-identical to the tested tree (git diff vs the saved patch: identical; the new helper file compared byte-for-byte). The numbers above come from that run: 33 of 35 new tests fail against the unfixed source with the exact errors quoted in the per-item sections.

Builds

src/AiDotNet.csproj compiles with 0 errors for net10.0, net8.0 and net471 (the library multi-targets all three; net471 lacks several modern BCL APIs, so it is built explicitly).

---

Adversarial review

Default parameter counts (measured, before → after): VideoMAE 65,788,816 → 65,788,816 · TimeSformer 113,541,120 → 113,541,120 · VideoCLIP 211,126,273 → 211,126,273 (would have been 38,004,480 without the walk fix) · RAFT 6,277,122 → 6,277,122 · SKFlow 300,098 → 300,098 · ProPainter 22,008,067 → 22,008,067 · BasicVSR++ 0 → 0 (fully lazy).

Clone / serialize:

• VideoMAE — clone preserves parameters, NumFrames and NormalizeTarget; the clone can run PretrainMAE; serialize → deserialize preserves parameters, NormalizeTarget, and produces bit-identical predictions.
• TimeSformer (16-frame architecture) — clone preserves parameters and now NumFrames = 16 (before: 16 → 8), and deserialize now yields an identical prediction (before: different, because the rebuilt model's frame count came back as 8 and re-grouped the tokens).
• VideoCLIP — clone preserves parameters and NumFrames.
• SKFlow — clone predicts the same unbatched [2,16,16] (before: the clone threw the rank-4 ArgumentException).

---

Blast radius / file map

Source (10 files, all under src/Video)

File	Change
Video/OpticalFlowBase.cs	rank-3 promotion + squeeze in PredictCore; body split into PredictBatchedPair; both [TensorLayout]s BatchOptional; OutputAxesFor(3)
Video/Motion/RAFT.cs	rank-3 promotion + squeeze; explicit rank guard
Video/ActionRecognition/TimeSformer.cs	frame count from the architecture; unbatched squeeze; per-clip softmax; ctor doc-comment placement
Video/ActionRecognition/VideoMAE.cs	pretraining step; sparse masking; decoder GELU removed; tape-built objective + normalised target; delegates frame resolution
Video/Options/VideoMAEOptions.cs	new NormalizeTarget option (default true)
Video/VideoClipFrameCount.cs	new: the one frame-count rule, internal
Video/Enhancement/BasicVSRPlusPlus.cs	one-frame clip in EnhanceVideo
Video/Generation/OpenSora.cs	rank-3 latent in PredictCore
Video/Inpainting/ProPainter.cs	rank-3 frame in PredictCore
Video/Understanding/VideoCLIP.cs	declared clip shape; frame count from the architecture; per-frame TryGetArchitectureInputShape

Nothing outside src/Video is touched — no change to NeuralNetworkBase, LayerHelper, the generators or any shared layer. No visibility was widened beyond one private → internal inside VideoMAE for its own test (the same pattern the base branch already uses for DecodeForReconstruction, CreateTubeMask and ComputeReconstructionLoss).

Tests (5 files): new UnitTests/Video/Motion/OpticalFlowDeclaredInputTests.cs (15), UnitTests/Video/VideoModelDeclaredInputTests.cs (7), UnitTests/Video/ActionRecognition/TimeSformerFrameCountTests.cs (7), UnitTests/Video/ActionRecognition/VideoMAEPretrainingFidelityTests.cs (6); plus VideoMAEPretrainingTests.cs — one existing exact-value test now constructs the model with NormalizeTarget = false, because it pins the raw-pixel correspondence by hand; the normalised default is pinned by a new test.

Behavioural changes callers could notice (all of them the point of the PR): unbatched optical-flow Predict now returns [2,H,W] instead of throwing; TimeSformer's unbatched Predict returns [classes] instead of [1, classes]; a batched TimeSformer Predict returns per-clip probabilities instead of batch-normalised ones; PretrainMAE now mutates weights and its loss is on normalised targets by default; a 16-frame TimeSformer / VideoCLIP architecture now builds a 16-/N-frame model instead of silently using the default.

---

What this does not do

Found while working, not fixed here — each needs its own change and its own test:

• VideoMAE.Train calls TrainWithTape(input, expectedOutput) without the constructor's _optimizer, so a user-supplied optimizer is silently ignored (604 other call sites pass it; TimeSformer does). PretrainMAE in this PR does use it. Left out because a regression test needs a non-fusable optimizer probe to observe which optimizer actually ran.
• OpenSora's [TensorLayout(Input)] is [Batch, Time, Features] while PredictCore takes [B, C, H, W], and its Train computes a gradient it never back-propagates. Its default config is ~446M parameters, so the declared-shape probe is only practical at a small config.
• ProPainter's inference always applies an all-zero mask, so Predict never actually inpaints anything.
• BasicVSR++ reads rank 4 as [F, C, H, W] while the inherited super-resolution layout says [B, C, H, W]; and VideoSuperResolutionBase.EstimateFlow defaults to an untrained full-size new RAFT<T>(), so every VSR model that does not override it aligns frames with random flow.
• GMFlow's rank-3 pair path splits Shape[0] / 2 with no even-channel check.

Plus the four declared-shape models in the decision section above.

🤖 Generated with Claude Code

https://claude.ai/code/session_017otuSGr3GdmvPoYLbiaWR2

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated
aidotnet-playground-api	Ready	Preview	Sep 11, 2026 10:11pm UTC

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (2)

• master
• integration/.*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Advanced

Run ID: 8101c0e1-742a-40f2-88d0-3903ed1f6bfb

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}