Arbitrary-Rank Block Masks

Masked prediction is not intrinsically an image operation. A model may hide intervals in a signal, rectangles in an image, cuboids in a volume, or blocks in a higher-dimensional simulation field. This module therefore describes a mask by two rank-indexed vectors:

• shape : Vector Nat d gives the tensor extents;
• blocks : Vector (Option Nat) d selects the axes that form a block grid.

none means that an axis does not participate in the block index. some k groups that axis into consecutive blocks of width k. The selected block-grid coordinates are flattened in row-major order, and one congruence class modulo period is hidden. A zero period, a zero block width, a rank mismatch, or an out-of-bounds coordinate hides nothing.

The executable mask and its coordinate theorems use the same finite predicate. Consequently the runtime training sample cannot silently use a different patch convention from the one stated in Lean.

def TorchLean.ssl.blockIndexAux : List ℕ → List (Option ℕ) → List ℕ → ℕ → Bool → Option ℕ

Compute the row-major index of the block containing coordinate.

Axes marked none are ignored, so the same spatial mask is repeated across batch, channel, token, or feature axes. The final Boolean records whether at least one axis participates in the block grid.

def TorchLean.ssl.blockIndex {d : ℕ} (shape : Vector ℕ d) (blocks : Vector (Option ℕ) d) (coordinate : Vector ℕ d) : Option ℕ

Row-major block index, or none for an invalid/degenerate block description.

def TorchLean.ssl.blockHidden {d : ℕ} (shape : Vector ℕ d) (blocks : Vector (Option ℕ) d) (period offset : ℕ) (coordinate : Vector ℕ d) : Bool

Whether a coordinate belongs to the selected congruence class of blocks.

def TorchLean.ssl.scalarAt {α : Type} (dims : List ℕ) : Spec.Tensor α (Spec.Shape.ofList dims) → List ℕ → Option α

Read a scalar from a shape-indexed tensor using runtime coordinates.

def TorchLean.ssl.blockHiddenList (shape : List ℕ) (blocks : List (Option ℕ)) (period offset : ℕ) (coordinate : List ℕ) : Bool

List-level hidden-coordinate predicate used by the recursive tensor implementation.

def TorchLean.ssl.blockMaskAux (shape : List ℕ) (blocks : List (Option ℕ)) (period offset : ℕ) (coordinatePrefix dims : List ℕ) : Spec.Tensor Float (Spec.Shape.ofList dims) → Spec.Tensor Float (Spec.Shape.ofList dims)

Recursively apply a block mask while accumulating the current coordinate.

def TorchLean.ssl.blockMask {d : ℕ} (shape : Vector ℕ d) (blocks : Vector (Option ℕ) d) (period offset : ℕ) (x : Spec.Tensor Float (Spec.Shape.ofList shape.toList)) : Spec.Tensor Float (Spec.Shape.ofList shape.toList)

Set every scalar in a selected block to zero, preserving the tensor's arbitrary-rank shape.

For example, policies [none, some 4, some 4] repeat a 4-by-4 block mask across the first axis; [some 8] masks intervals in a signal; and [some 2, some 2, some 2] masks volume blocks.

theorem TorchLean.ssl.scalarAt_blockMaskAux (shape : List ℕ) (blocks : List (Option ℕ)) (period offset : ℕ) (coordinatePrefix dims : List ℕ) (x : Spec.Tensor Float (Spec.Shape.ofList dims)) (coordinates : List ℕ) : scalarAt dims (blockMaskAux shape blocks period offset coordinatePrefix dims x) coordinates = Option.map (fun (value : Float) => if blockHiddenList shape blocks period offset (coordinatePrefix ++ coordinates) = true then 0.0 else value) (scalarAt dims x coordinates)

theorem TorchLean.ssl.blockMask_scalarAt {d : ℕ} (shape : Vector ℕ d) (blocks : Vector (Option ℕ) d) (period offset : ℕ) (x : Spec.Tensor Float (Spec.Shape.ofList shape.toList)) (coordinate : Vector ℕ d) : scalarAt shape.toList (blockMask shape blocks period offset x) coordinate.toList = Option.map (fun (value : Float) => if blockHidden shape blocks period offset coordinate = true then 0.0 else value) (scalarAt shape.toList x coordinate.toList)

Exact coordinate semantics of blockMask, including out-of-bounds coordinates.

theorem TorchLean.ssl.blockMask_hidden_scalar_eq_zero {d : ℕ} (shape : Vector ℕ d) (blocks : Vector (Option ℕ) d) (period offset : ℕ) (x : Spec.Tensor Float (Spec.Shape.ofList shape.toList)) (coordinate : Vector ℕ d) (value : Float) (hValue : scalarAt shape.toList x coordinate.toList = some value) (hHidden : blockHidden shape blocks period offset coordinate = true) : scalarAt shape.toList (blockMask shape blocks period offset x) coordinate.toList = some 0.0

A selected in-bounds coordinate is exactly zero after masking.

theorem TorchLean.ssl.blockMask_visible_scalar_eq_input {d : ℕ} (shape : Vector ℕ d) (blocks : Vector (Option ℕ) d) (period offset : ℕ) (x : Spec.Tensor Float (Spec.Shape.ofList shape.toList)) (coordinate : Vector ℕ d) (value : Float) (hValue : scalarAt shape.toList x coordinate.toList = some value) (hVisible : blockHidden shape blocks period offset coordinate = false) : scalarAt shape.toList (blockMask shape blocks period offset x) coordinate.toList = some value

A visible in-bounds coordinate is copied unchanged by the mask.

def TorchLean.ssl.blockMaskBatch {d : ℕ} (batch : ℕ) (shape : Vector ℕ d) (blocks : Vector (Option ℕ) d) (period offset : ℕ) (x : Spec.Tensor Float (Spec.Shape.dim batch (Spec.Shape.ofList shape.toList))) : Spec.Tensor Float (Spec.Shape.dim batch (Spec.Shape.ofList shape.toList))

Apply the same block mask independently to each row of a batch.

theorem TorchLean.ssl.blockMaskBatch_scalarAt {d : ℕ} (batch : ℕ) (shape : Vector ℕ d) (blocks : Vector (Option ℕ) d) (period offset : ℕ) (x : Spec.Tensor Float (Spec.Shape.dim batch (Spec.Shape.ofList shape.toList))) (row : Fin batch) (coordinate : Vector ℕ d) : scalarAt shape.toList (Spec.get (blockMaskBatch batch shape blocks period offset x) row) coordinate.toList = Option.map (fun (value : Float) => if blockHidden shape blocks period offset coordinate = true then 0.0 else value) (scalarAt shape.toList (Spec.get x row) coordinate.toList)

Coordinate semantics of one row of blockMaskBatch.

def TorchLean.ssl.blockMaeSample {d : ℕ} (batch reconDim : ℕ) (shape : Vector ℕ d) (blocks : Vector (Option ℕ) d) (period offset : ℕ) (hRecon : reconDim ≤ (Spec.Shape.ofList shape.toList).size) (x : Spec.Tensor Float (Spec.Shape.dim batch (Spec.Shape.ofList shape.toList))) : Sample.Supervised Float (Spec.Shape.dim batch (Spec.Shape.ofList shape.toList)) (Spec.Shape.dim batch (Spec.Shape.dim reconDim Shape.scalar))

Create a masked-reconstruction sample from a batch of arbitrary-rank tensors.

The model input retains its original shape. The target is a row-major prefix of the unmasked source because TorchLean's compact decoder heads produce matrices; reconDim may be the entire sample or a smaller prefix for an experiment.

theorem TorchLean.ssl.blockMaeSample_input_eq_mask {d : ℕ} (batch reconDim : ℕ) (shape : Vector ℕ d) (blocks : Vector (Option ℕ) d) (period offset : ℕ) (hRecon : reconDim ≤ (Spec.Shape.ofList shape.toList).size) (x : Spec.Tensor Float (Spec.Shape.dim batch (Spec.Shape.ofList shape.toList))) : Sample.x (blockMaeSample batch reconDim shape blocks period offset hRecon x) = blockMaskBatch batch shape blocks period offset x

The model input of a block-MAE sample is exactly the masked source batch.

theorem TorchLean.ssl.blockMaeSample_target_eq_source_prefix {d : ℕ} (batch reconDim : ℕ) (shape : Vector ℕ d) (blocks : Vector (Option ℕ) d) (period offset : ℕ) (hRecon : reconDim ≤ (Spec.Shape.ofList shape.toList).size) (x : Spec.Tensor Float (Spec.Shape.dim batch (Spec.Shape.ofList shape.toList))) : Sample.y (blockMaeSample batch reconDim shape blocks period offset hRecon x) = Tensor.flattenBatchPrefix batch reconDim hRecon x

The target of a block-MAE sample is the requested prefix of the unmasked source batch.

def TorchLean.ssl.blockMaeReconstructionIndices (reconDim : ℕ) : List (Fin reconDim)

Every decoder coordinate participates in the compact reconstruction objective.

def TorchLean.ssl.blockMaeRowPredictiveContract {d : ℕ} (batch reconDim : ℕ) (shape : Vector ℕ d) (blocks : Vector (Option ℕ) d) (period offset : ℕ) (hRecon : reconDim ≤ (Spec.Shape.ofList shape.toList).size) (x : Spec.Tensor Float (Spec.Shape.dim batch (Spec.Shape.ofList shape.toList))) (prediction : Spec.Tensor Float (Spec.Shape.dim batch (Spec.Shape.dim reconDim Shape.scalar))) (row : Fin batch) (loss : Float → Float → ℕ) : NN.MLTheory.SelfSupervised.PredictiveViewContract reconDim Unit Float Float Float

One batch row of block-MAE training as a finite predictive-view contract.

theorem TorchLean.ssl.blockMaeRow_predictive_objective_eq_maeLoss {d : ℕ} (batch reconDim : ℕ) (shape : Vector ℕ d) (blocks : Vector (Option ℕ) d) (period offset : ℕ) (hRecon : reconDim ≤ (Spec.Shape.ofList shape.toList).size) (x : Spec.Tensor Float (Spec.Shape.dim batch (Spec.Shape.ofList shape.toList))) (prediction : Spec.Tensor Float (Spec.Shape.dim batch (Spec.Shape.dim reconDim Shape.scalar))) (row : Fin batch) (loss : Float → Float → ℕ) : NN.MLTheory.SelfSupervised.predictiveViewObjective (blockMaeRowPredictiveContract batch reconDim shape blocks period offset hRecon x prediction row loss) = NN.MLTheory.SelfSupervised.maeLoss (blockMaeReconstructionIndices reconDim) (matrixRowAsPatchBatch batch reconDim (Sample.y (blockMaeSample batch reconDim shape blocks period offset hRecon x)) row) (matrixRowAsPrediction batch reconDim prediction row) loss

The runnable block-MAE row objective is exactly the finite MAE objective.