Fix docs for TransformerBlock and remove incompatible functions

codecov.yml

@@ -0,0 +1,3 @@
+ignore:
+  - "src/convolution/UNet"
+

src/Onion.jl

@@ -20,9 +20,7 @@ using .Utils
 export glut
 export like, zeros_like, ones_like, falses_like, trues_like
 export watmul, ⨝
-export self_att_padding_mask
-export cross_att_padding_mask
-export causal_mask
+export causal_mask, self_att_padding_mask, cross_att_padding_mask
 export bf16
 
 const Maybe{T} = Union{T,Nothing}

src/Utils/Utils.jl

@@ -28,9 +28,7 @@ include("ofeltype.jl")
 export ofeltype
 
 include("masks.jl")
-export self_att_padding_mask
-export cross_att_padding_mask
-export causal_mask
+export causal_mask, self_att_padding_mask, cross_att_padding_mask
 
 include("b16.jl")
 export bf16

src/Utils/masks.jl

@@ -62,7 +62,6 @@ function cross_att_padding_mask(padmask, other_dim; T=Float32)
     return log.(repeat(pm, einops"n ... -> n m ..."; m=other_dim))
 end
 
-
 function causal_mask(x::AbstractArray{<:AbstractFloat})
     n = size(x, 2)
     mask = like(-Inf, x, n, n)

src/convolution/UNet/FlexibleUNet.jl

@@ -2,7 +2,6 @@
     FlexibleUNet(;
         in_channels=3,
         out_channels=3,
-        depth=3,
         base_channels=64,
         channel_multipliers=[1, 2, 4],
         time_embedding=false,
@@ -15,29 +14,28 @@
     )
 
 A flexible UNet architecture with configurable depth and channel dimensions.
+The depth is determined by the length of `channel_multipliers`.
 Supports optional time and class embeddings for diffusion models and conditional generation.
 
 # Arguments
 - `in_channels=3`: Number of input channels
 - `out_channels=3`: Number of output channels
-- `depth=3`: Number of encoder/decoder blocks
 - `base_channels=64`: Base channel dimension (multiplied at each level)
-- `channel_multipliers=[1, 2, 4]`: Multipliers for channel dimensions at each level
+- `channel_multipliers=[1, 2, 4]`: Multipliers for channel dimensions at each level. The length determines the depth (number of encoder/decoder blocks)
 - `time_embedding=false`: Whether to use time embeddings
 - `num_classes=0`: Number of class labels for conditional generation
 - `embedding_dim=128`: Dimension for class embeddings
 - `time_emb_dim=256`: Dimension for time embeddings
 - `dropout=0.0`: Dropout probability to apply to inner layers
-- `dropout_depth=0`: Number of layers to apply dropout to, starting from the innermost layers (0 means no dropout). Maximum value is 1+depth (bottleneck + all encoding/decoding levels)
+- `dropout_depth=0`: Number of layers to apply dropout to, starting from the innermost layers (0 means no dropout). Maximum value is 1+length(channel_multipliers) (bottleneck + all encoding/decoding levels)
 - `activation=relu`: Activation function to use throughout the network
 
 # Examples
 ```julia
-# Basic model without dropout
+# Basic model without dropout (depth=4 from channel_multipliers length)
 model = Onion.UNet.FlexibleUNet(
     in_channels=3,
     out_channels=3,
-    depth=4,
     base_channels=32,
     channel_multipliers=[1, 2, 4, 8],
     time_embedding=true
@@ -47,10 +45,10 @@ model = Onion.UNet.FlexibleUNet(
 model = Onion.UNet.FlexibleUNet(
     in_channels=3,
     out_channels=3,
-    depth=4,
     base_channels=32,
     channel_multipliers=[1, 2, 4, 8],
     time_embedding=true,
+    num_classes=10,
     dropout=0.2,
     dropout_depth=3
 )
@@ -67,14 +65,17 @@ struct FlexibleUNet{E,B,D,FC,T}
     decoders::D
     final_conv::FC
     time_embed::T
+    in_channels::Int
+    out_channels::Int
+    time_embedding::Bool
+    num_classes::Int
 end
 
 @layer FlexibleUNet
 
 function FlexibleUNet(;
     in_channels=3,
     out_channels=3,
-    depth=3,
     base_channels=64,
     channel_multipliers=[1, 2, 4], # Multipliers for each level
     time_embedding=false,
@@ -85,15 +86,8 @@ function FlexibleUNet(;
     dropout_depth=0,
     activation=relu
 )
-    # Ensure we have enough channel multipliers for the requested depth
-    if length(channel_multipliers) < depth
-        # Extend with the last multiplier (create new array, don't mutate)
-        channel_multipliers = vcat(channel_multipliers,
-                fill(channel_multipliers[end], depth - length(channel_multipliers)))
-    elseif length(channel_multipliers) > depth
-        # Trim to the requested depth (create new array, don't mutate)
-        channel_multipliers = channel_multipliers[1:depth]
-    end
+    # Depth is determined by the length of channel_multipliers
+    depth = length(channel_multipliers)
 
     # Calculate actual channel numbers
     channels = [base_channels * m for m in channel_multipliers]
@@ -160,7 +154,8 @@ function FlexibleUNet(;
     encoders_tuple = Tuple(encoders)
     decoders_tuple = Tuple(decoders)
 
-    FlexibleUNet(encoders_tuple, bottleneck, decoders_tuple, final_conv, time_embed)
+    FlexibleUNet(encoders_tuple, bottleneck, decoders_tuple, final_conv, time_embed, 
+                 in_channels, out_channels, time_embedding, num_classes)
 end
 
 # Process encoders and collect skip connections without mutations
@@ -221,6 +216,9 @@ end
 
 # Standard forward pass without time embedding - using foldl to avoid mutations
 function (model::FlexibleUNet)(x)
+    @assert ndims(x) >= 3 "Input must be at least 3D (height, width, channels) or 4D (height, width, channels, batch)"
+    @assert size(x, 3) == model.in_channels "Input channels mismatch: expected $(model.in_channels), got $(size(x, 3))"
+
     # Apply encoder blocks and collect skip connections
     x, skip_connections = process_encoders(x, model.encoders)
 
@@ -237,6 +235,9 @@ end
 
 # Forward pass with time embedding - using foldl to avoid mutations
 function (model::FlexibleUNet)(x, t::T) where T <: AbstractArray
+    @assert ndims(x) >= 3 "Input must be at least 3D (height, width, channels) or 4D (height, width, channels, batch)"
+    @assert size(x, 3) == model.in_channels "Input channels mismatch: expected $(model.in_channels), got $(size(x, 3))"
+
     t = model.time_embed(t)
 
     # Apply encoder blocks and collect skip connections
@@ -255,6 +256,10 @@ end
 
 # Forward pass with time embedding and class labels - using foldl to avoid mutations
 function (model::FlexibleUNet)(x, t::T, labels::L) where {T <: AbstractArray, L <: AbstractArray}
+    @assert ndims(x) >= 3 "Input must be at least 3D (height, width, channels) or 4D (height, width, channels, batch)"
+    @assert size(x, 3) == model.in_channels "Input channels mismatch: expected $(model.in_channels), got $(size(x, 3))"
+    @assert model.num_classes > 0 "Model was constructed with num_classes=0, but labels were provided. Set num_classes > 0 when creating the model."
+
     t = model.time_embed(t, labels)
 
     # Apply encoder blocks and collect skip connections

src/norm/AdaLN.jl

@@ -1,7 +1,12 @@
 """
     AdaLN(dim::Int, cond_dim::Int)
 
-Adaptive Layer Normalization.
+Adaptive Layer Normalization. 
+
+Adaptive Layer Normalization modulates the normalized output using learned scale and shift parameters computed from a conditioning tensor.
+
+The forward pass computes: `output = normalized(x) * (1 + scale(cond)) + shift(cond)`
+where `scale` and `shift` are learned Dense layers that transform the conditioning tensor.
 
 ```julia
 aln = AdaLN(5, 3)

src/positional-encoding/RoPE.jl

@@ -13,7 +13,7 @@ t = TransformerBlock(dim, n_heads, n_kv_heads)
 h = randn(Float32, dim, seqlen, 1)
 
 rope = RoPE(dim ÷ n_heads, 1000)
-h = t(h, 1, rope[1:seqlen]) #Note the subsetting to match seqlen
+h = t(h; rope=rope[1:seqlen]) #Note the subsetting to match seqlen
 ```
 """
 struct RoPE{A<:AbstractArray}

src/transformers/block.jl

@@ -3,6 +3,44 @@
 
 Transformer block for GQAttention (as in Llama3).
 
+# Constructor
+
+Creates a transformer block with the specified dimensions:
+
+- `dim`: Model dimension
+- `n_heads`: Number of attention heads
+- `n_kv_heads`: Number of key-value heads (defaults to `n_heads`)
+- `ff_hidden_dim`: Feed-forward hidden dimension (defaults to `4 * dim`)
+- `norm_eps`: Normalization epsilon (defaults to `1f-5`)
+- `qkv_bias`: Whether to use bias in QKV projections (defaults to `false`)
+
+# Forward Pass
+
+Call the block as a function to perform the forward pass:
+
+```julia
+(block::TransformerBlock)(x, xs...; cond=nothing, pair_feats=nothing, pair=block.pair_proj(pair_feats), kws...)
+```
+
+**Parameters:**
+
+- `x`: Input tensor of shape `(dim, seqlen, batch)`
+- `xs...`: Additional positional arguments passed to attention (e.g., separate key/value tensors for cross-attention)
+- `cond`: Optional conditioning tensor for adaptive normalization. When provided, this is passed to both `attention_norm` and `ffn_norm`. 
+  - For regular `RMSNorm` or `LayerNorm`, this parameter is ignored (can be `nothing`)
+  - For `AdaLN` (Adaptive Layer Normalization), this is required and should have shape `(cond_dim, batch)` where `cond_dim` matches the dimension specified when creating `AdaLN`
+  - When using `AdaLN`, the conditioning tensor is passed directly to the normalization layer. `AdaLN` uses learned linear transformations (Dense layers) to compute scale and shift parameters from the conditioning tensor, which then modulate the normalized output: `output = normalized(x) * (1 + scale(cond)) + shift(cond)`
+- `pair_feats`: Optional pair features for attention
+- `kws...`: Additional keyword arguments passed to attention, including:
+  - `rope`: Optional RoPE function or callable for query positions (e.g., `rope=rope[1:seqlen]`)
+  - `krope`: Optional RoPE function or callable for key positions (defaults to `rope`)
+  - `causal::Bool=false`: Enable causal masking
+  - `kpad_mask`: Optional padding mask for keys in **probability space** (values in `[0, 1]` where `1` indicates valid key position and `0` indicates padded). Should have shape `(kl, batch)` where `kl` is key length. Pass a sequence-level mask directly (e.g., `ones(Float32, seqlen, batch)` with padded positions set to `0`). The mask is automatically converted to log-space and broadcast over query length and heads by `apply_pad_mask`.
+
+## Examples
+
+### Basic forward pass
+
 ```julia
 dim = 64
 n_heads = 8
@@ -11,15 +49,151 @@ seqlen = 10
 
 rope = RoPE(dim ÷ n_heads, 1000)
 t = TransformerBlock(dim, n_heads, n_kv_heads)
+h = randn(Float32, dim, seqlen, 1)
+
+# Forward pass
+h = t(h; rope=rope[1:seqlen])
+```
+
+### With causal masking
 
+Use `causal=true` to enable causal masking:
+
+```julia
+dim = 64
+n_heads = 8
+seqlen = 10
+
+rope = RoPE(dim ÷ n_heads, 1000)
+t = TransformerBlock(dim, n_heads)
 h = randn(Float32, dim, seqlen, 1)
 
-#Use without a mask:
-h = t(h, 1, rope[1:seqlen])
+# Forward pass with causal mask
+h = t(h; rope=rope[1:seqlen], causal=true)
+```
+
+### With padding mask
+
+Use `kpad_mask` for self-attention padding. Pass a sequence-level mask in **probability space** (values in `[0, 1]` where `1` indicates valid position and `0` indicates padding):
+
+```julia
+dim = 64
+n_heads = 8
+seqlen = 10
+batch = 2
+
+rope = RoPE(dim ÷ n_heads, 1000)
+t = TransformerBlock(dim, n_heads)
+h = randn(Float32, dim, seqlen, batch)
+
+# Create padding mask: sequence-level mask where 1 indicates valid position, 0 indicates padding
+# Example: batch 1 has all positions valid, batch 2 has position 10 padded
+# Shape should be (seqlen, batch)
+kpad_mask = ones(Float32, seqlen, batch)
+kpad_mask[10, 2] = 0  # Mark position 10 in batch 2 as padding
+
+# Forward pass with padding mask
+# Note: The mask is automatically converted to log-space and broadcast over query length and heads
+h = t(h; rope=rope[1:seqlen], kpad_mask=kpad_mask)
+```
+
+### With cross-attention padding mask
+
+Use `kpad_mask` when using TransformerBlock with different query and key sequences. Pass a key-level mask in **probability space** (values in `[0, 1]` where `1` indicates valid key position and `0` indicates padded):
+
+```julia
+dim = 64
+n_heads = 8
+q_seqlen = 10
+k_seqlen = 12
+batch = 2
+
+rope = RoPE(dim ÷ n_heads, 1000)
+t = TransformerBlock(dim, n_heads)
+q = randn(Float32, dim, q_seqlen, batch)
+k = randn(Float32, dim, k_seqlen, batch)
+
+# Create padding mask for keys: sequence-level mask where 1 indicates valid position
+# Shape should be (k_seqlen, batch)
+kpad_mask = ones(Float32, k_seqlen, batch)
+kpad_mask[12, 2] = 0  # Mark position 12 in batch 2 as padding
+
+# Forward pass: pass key as positional argument after query
+# Note: The mask is automatically converted to log-space and broadcast over query length and heads
+h = t(q, k; rope=rope[1:q_seqlen], krope=rope[1:k_seqlen], kpad_mask=kpad_mask)
+```
+
+### Combining causal and padding masks
+
+You can combine both causal and padding masks. Padding masks should be in **probability space**:
+
+```julia
+dim = 64
+n_heads = 8
+seqlen = 10
+batch = 2
+
+rope = RoPE(dim ÷ n_heads, 1000)
+t = TransformerBlock(dim, n_heads)
+h = randn(Float32, dim, seqlen, batch)
+
+# Create padding mask: shape should be (seqlen, batch)
+kpad_mask = ones(Float32, seqlen, batch)
+kpad_mask[10, 2] = 0  # Mark position 10 in batch 2 as padding
+
+# Forward pass with both causal and padding masks
+# Note: The padding mask is automatically converted to log-space and broadcast over query length and heads
+h = t(h; rope=rope[1:seqlen], causal=true, kpad_mask=kpad_mask)
+```
+
+### With adaptive normalization (AdaLN)
+
+When using `AdaLN` (Adaptive Layer Normalization) instead of `RMSNorm`, you must pass a conditioning tensor via the `cond` parameter. 
+`AdaLN` takes the conditioning tensor and passes it through learned linear transformations (Dense layers) to compute scale and shift parameters. 
+These learned parameters then modulate the normalized output according to: `output = normalized(x) * (1 + scale(cond)) + shift(cond)`.
+
+The conditioning tensor is passed directly to the normalization layers (`attention_norm` and `ffn_norm`) during the forward pass.
+
+Use `AdaTransformerBlock` for a convenience constructor that creates a TransformerBlock with `AdaLN`:
+
+```julia
+dim = 64
+cond_dim = 32  # Dimension of conditioning tensor
+n_heads = 8
+seqlen = 10
+batch = 2
+
+rope = RoPE(dim ÷ n_heads, 1000)
+t = AdaTransformerBlock(dim, cond_dim, n_heads)
+h = randn(Float32, dim, seqlen, batch)
+
+# Create conditioning tensor
+cond = randn(Float32, cond_dim, batch)
+
+# Forward pass with conditioning tensor
+h = t(h; rope=rope[1:seqlen], cond=cond)
+
+# You can still use masks with adaptive normalization
+h = t(h; rope=rope[1:seqlen], cond=cond, causal=true)
+```
+
+Alternatively, you can manually create a TransformerBlock with AdaLN:
+
+```julia
+using Onion: AdaLN
+
+dim = 64
+cond_dim = 32
+n_heads = 8
+t = TransformerBlock(
+    dim, n_heads;
+    attention_norm = AdaLN(dim, cond_dim),
+    ffn_norm = AdaLN(dim, cond_dim)
+)
 
-#Use with a causal mask:
-mask = Onion.causal_mask(h)
-h = t(h, 1, rope[1:seqlen], mask)
+# Usage is the same
+cond = randn(Float32, cond_dim, batch)
+h = t(h; rope=rope[1:seqlen], cond=cond)
 ```
 """
 @concrete struct TransformerBlock
@@ -49,6 +223,7 @@ function TransformerBlock(
     )
 end
 
+# Hidden from docs - forward pass documentation is in the constructor docstring above
 function (block::TransformerBlock)(
     x, xs...;
     cond=nothing, pair_feats=nothing,