chickadee » nanograd » make-flatten

Creates a flatten layer that converts multi-dimensional tensors to 1D or 2D.

Input shapes and outputs:

• 4D: (N, C, H, W) → (N, C*H*W)
• 3D: (C, H, W) → (C*H*W)
• 2D: (N, features) → (N, features) (no change)
• 1D: (features,) → (features,) (no change)

(define flatten (make-flatten name: "Flatten"))

; Flatten batch of feature maps
(define features (make-tensor32 data '(32 64 8 8)))
(define flattened (forward flatten features))  ; Shape: (32, 4096)