-module(mlx_nn). %% Advanced neural network layers and architectures for SOTA ML -export([ %% Core layer types linear/3, linear/4, embedding/3, embedding/4, %% Convolutional layers conv1d/4, conv1d/5, conv1d/6, conv2d/4, conv2d/5, conv2d/6, conv3d/4, conv3d/5, conv3d/6, conv_transpose1d/4, conv_transpose2d/4, conv_transpose3d/4, %% Pooling layers max_pool1d/2, max_pool1d/3, max_pool2d/2, max_pool2d/3, avg_pool1d/2, avg_pool1d/3, avg_pool2d/2, avg_pool2d/3, adaptive_avg_pool1d/2, adaptive_avg_pool2d/2, global_avg_pool/1, global_max_pool/1, %% Attention mechanisms scaled_dot_product_attention/3, scaled_dot_product_attention/4, multi_head_attention/6, multi_head_attention/7, self_attention/3, self_attention/4, cross_attention/5, cross_attention/6, %% Transformer components transformer_encoder_layer/6, transformer_encoder_layer/7, transformer_decoder_layer/7, transformer_decoder_layer/8, positional_encoding/3, positional_encoding/4, %% Normalization layers layer_norm/2, layer_norm/3, layer_norm/4, batch_norm/2, batch_norm/3, batch_norm/4, group_norm/3, group_norm/4, instance_norm/2, instance_norm/3, rms_norm/2, rms_norm/3, %% Dropout and regularization dropout/2, dropout/3, dropout2d/2, dropout3d/2, alpha_dropout/2, feature_alpha_dropout/2, %% Advanced architectures residual_block/3, residual_block/4, dense_block/3, dense_block/4, squeeze_excitation/2, squeeze_excitation/3, %% Recurrent layers lstm_cell/4, lstm_cell/5, gru_cell/3, gru_cell/4, rnn_cell/3, rnn_cell/4, %% Modern architectures vision_transformer_block/4, vision_transformer_block/5, swin_transformer_block/5, swin_transformer_block/6, efficientnet_block/5, efficientnet_block/6, %% Initialization functions xavier_uniform/1, xavier_normal/1, kaiming_uniform/1, kaiming_normal/1, orthogonal_init/1, identity_init/1 ]). %% Core layer implementations linear(Input, Weight, Bias) -> linear(Input, Weight, Bias, #{}). linear(Input, Weight, Bias, Options) -> % Linear transformation: y = xW^T + b {ok, Output} = mlx:matmul(Input, mlx:transpose(Weight)), case Bias of undefined -> {ok, Output}; _ -> mlx:add(Output, Bias) end. embedding(Input, EmbeddingMatrix, VocabSize) -> embedding(Input, EmbeddingMatrix, VocabSize, #{}). embedding(Input, EmbeddingMatrix, VocabSize, Options) -> % Embedding layer: lookup embeddings by index mlx_nif:embedding(Input, EmbeddingMatrix, VocabSize, Options). %% Convolutional layers conv1d(Input, Weight, Bias, Stride) -> conv1d(Input, Weight, Bias, Stride, 0, 1). conv1d(Input, Weight, Bias, Stride, Padding) -> conv1d(Input, Weight, Bias, Stride, Padding, 1). conv1d(Input, Weight, Bias, Stride, Padding, Dilation) -> {ok, Output} = mlx_nif:conv1d(Input, Weight, #{ stride => Stride, padding => Padding, dilation => Dilation }), case Bias of undefined -> {ok, Output}; _ -> mlx:add(Output, Bias) end. conv2d(Input, Weight, Bias, Stride) -> conv2d(Input, Weight, Bias, Stride, 0, 1). conv2d(Input, Weight, Bias, Stride, Padding) -> conv2d(Input, Weight, Bias, Stride, Padding, 1). conv2d(Input, Weight, Bias, Stride, Padding, Dilation) -> {ok, Output} = mlx_nif:conv2d(Input, Weight, #{ stride => Stride, padding => Padding, dilation => Dilation }), case Bias of undefined -> {ok, Output}; _ -> mlx:add(Output, Bias) end. conv3d(Input, Weight, Bias, Stride) -> conv3d(Input, Weight, Bias, Stride, 0, 1). conv3d(Input, Weight, Bias, Stride, Padding) -> conv3d(Input, Weight, Bias, Stride, Padding, 1). conv3d(Input, Weight, Bias, Stride, Padding, Dilation) -> {ok, Output} = mlx_nif:conv3d(Input, Weight, #{ stride => Stride, padding => Padding, dilation => Dilation }), case Bias of undefined -> {ok, Output}; _ -> mlx:add(Output, Bias) end. conv_transpose1d(Input, Weight, Bias, Stride) -> mlx_nif:conv_transpose1d(Input, Weight, Bias, Stride). conv_transpose2d(Input, Weight, Bias, Stride) -> mlx_nif:conv_transpose2d(Input, Weight, Bias, Stride). conv_transpose3d(Input, Weight, Bias, Stride) -> mlx_nif:conv_transpose3d(Input, Weight, Bias, Stride). %% Pooling layers max_pool1d(Input, KernelSize) -> max_pool1d(Input, KernelSize, KernelSize). max_pool1d(Input, KernelSize, Stride) -> mlx_nif:max_pool1d(Input, KernelSize, Stride). max_pool2d(Input, KernelSize) -> max_pool2d(Input, KernelSize, KernelSize). max_pool2d(Input, KernelSize, Stride) -> mlx_nif:max_pool2d(Input, KernelSize, Stride). avg_pool1d(Input, KernelSize) -> avg_pool1d(Input, KernelSize, KernelSize). avg_pool1d(Input, KernelSize, Stride) -> mlx_nif:avg_pool1d(Input, KernelSize, Stride). avg_pool2d(Input, KernelSize) -> avg_pool2d(Input, KernelSize, KernelSize). avg_pool2d(Input, KernelSize, Stride) -> mlx_nif:avg_pool2d(Input, KernelSize, Stride). adaptive_avg_pool1d(Input, OutputSize) -> mlx_nif:adaptive_avg_pool1d(Input, OutputSize). adaptive_avg_pool2d(Input, OutputSize) -> mlx_nif:adaptive_avg_pool2d(Input, OutputSize). global_avg_pool(Input) -> % Global average pooling {ok, Shape} = mlx:shape(Input), SpatialDims = lists:seq(2, length(Shape) - 1), mlx:mean(Input, SpatialDims, true). global_max_pool(Input) -> % Global max pooling {ok, Shape} = mlx:shape(Input), SpatialDims = lists:seq(2, length(Shape) - 1), mlx:max(Input, SpatialDims, true). %% Attention mechanisms scaled_dot_product_attention(Query, Key, Value) -> scaled_dot_product_attention(Query, Key, Value, undefined). scaled_dot_product_attention(Query, Key, Value, Mask) -> % Scaled dot-product attention {ok, [_B, _N, DK]} = mlx:shape(Key), {ok, Scale} = mlx:array(1.0 / math:sqrt(DK)), % Compute attention scores {ok, Scores} = mlx:matmul(Query, mlx:transpose(Key, [0, 2, 1])), {ok, ScaledScores} = mlx:multiply(Scores, Scale), % Apply mask if provided ScoresWithMask = case Mask of undefined -> ScaledScores; _ -> {ok, MaskedScores} = mlx:where(Mask, ScaledScores, mlx:array(-1.0e9)), MaskedScores end, % Compute attention weights {ok, AttentionWeights} = mlx:softmax(ScoresWithMask, -1), % Apply attention to values {ok, Output} = mlx:matmul(AttentionWeights, Value), {ok, {Output, AttentionWeights}}. multi_head_attention(Query, Key, Value, WQ, WK, WV) -> multi_head_attention(Query, Key, Value, WQ, WK, WV, #{}). multi_head_attention(Query, Key, Value, WQ, WK, WV, Options) -> NumHeads = maps:get(num_heads, Options, 8), Dropout = maps:get(dropout, Options, 0.0), % Linear projections {ok, Q} = linear(Query, WQ, undefined), {ok, K} = linear(Key, WK, undefined), {ok, V} = linear(Value, WV, undefined), % Reshape for multi-head attention {ok, [BatchSize, SeqLen, DModel]} = mlx:shape(Q), DHead = DModel div NumHeads, {ok, QReshaped} = mlx:reshape(Q, [BatchSize, SeqLen, NumHeads, DHead]), {ok, KReshaped} = mlx:reshape(K, [BatchSize, SeqLen, NumHeads, DHead]), {ok, VReshaped} = mlx:reshape(V, [BatchSize, SeqLen, NumHeads, DHead]), % Transpose to [BatchSize, NumHeads, SeqLen, DHead] {ok, QTransposed} = mlx:transpose(QReshaped, [0, 2, 1, 3]), {ok, KTransposed} = mlx:transpose(KReshaped, [0, 2, 1, 3]), {ok, VTransposed} = mlx:transpose(VReshaped, [0, 2, 1, 3]), % Scaled dot-product attention {ok, {AttentionOutput, Weights}} = scaled_dot_product_attention( QTransposed, KTransposed, VTransposed), % Transpose back and reshape {ok, OutputTransposed} = mlx:transpose(AttentionOutput, [0, 2, 1, 3]), {ok, Output} = mlx:reshape(OutputTransposed, [BatchSize, SeqLen, DModel]), % Apply dropout if training case Dropout > 0.0 of true -> dropout(Output, Dropout); false -> {ok, Output} end. self_attention(Input, WQ, WK) -> self_attention(Input, WQ, WK, #{}). self_attention(Input, WQ, WK, Options) -> WV = maps:get(wv, Options, WK), multi_head_attention(Input, Input, Input, WQ, WK, WV, Options). cross_attention(Query, Key, Value, WQ, WK) -> cross_attention(Query, Key, Value, WQ, WK, #{}). cross_attention(Query, Key, Value, WQ, WK, Options) -> WV = maps:get(wv, Options, WK), multi_head_attention(Query, Key, Value, WQ, WK, WV, Options). %% Transformer components transformer_encoder_layer(Input, SelfAttnW, FFNWeights, LayerNorm1, LayerNorm2, Dropout) -> transformer_encoder_layer(Input, SelfAttnW, FFNWeights, LayerNorm1, LayerNorm2, Dropout, #{}). transformer_encoder_layer(Input, SelfAttnW, FFNWeights, LayerNorm1, LayerNorm2, Dropout, Options) -> % Self-attention block {ok, AttnOutput} = self_attention(Input, SelfAttnW, SelfAttnW, Options), {ok, AttnDropped} = dropout(AttnOutput, Dropout), {ok, AttnResidual} = mlx:add(Input, AttnDropped), {ok, AttnNormed} = layer_norm(AttnResidual, LayerNorm1), % Feed-forward block [W1, W2] = FFNWeights, {ok, FFNIntermediate} = linear(AttnNormed, W1, undefined), {ok, FFNActivated} = mlx:relu(FFNIntermediate), {ok, FFNOutput} = linear(FFNActivated, W2, undefined), {ok, FFNDropped} = dropout(FFNOutput, Dropout), {ok, FFNResidual} = mlx:add(AttnNormed, FFNDropped), layer_norm(FFNResidual, LayerNorm2). transformer_decoder_layer(Input, Memory, SelfAttnW, CrossAttnW, FFNWeights, LayerNorms, Dropout) -> transformer_decoder_layer(Input, Memory, SelfAttnW, CrossAttnW, FFNWeights, LayerNorms, Dropout, #{}). transformer_decoder_layer(Input, Memory, SelfAttnW, CrossAttnW, FFNWeights, LayerNorms, Dropout, Options) -> [LayerNorm1, LayerNorm2, LayerNorm3] = LayerNorms, % Self-attention block {ok, SelfAttnOutput} = self_attention(Input, SelfAttnW, SelfAttnW, Options), {ok, SelfAttnDropped} = dropout(SelfAttnOutput, Dropout), {ok, SelfAttnResidual} = mlx:add(Input, SelfAttnDropped), {ok, SelfAttnNormed} = layer_norm(SelfAttnResidual, LayerNorm1), % Cross-attention block {ok, CrossAttnOutput} = cross_attention(SelfAttnNormed, Memory, Memory, CrossAttnW, CrossAttnW, Options), {ok, CrossAttnDropped} = dropout(CrossAttnOutput, Dropout), {ok, CrossAttnResidual} = mlx:add(SelfAttnNormed, CrossAttnDropped), {ok, CrossAttnNormed} = layer_norm(CrossAttnResidual, LayerNorm2), % Feed-forward block [W1, W2] = FFNWeights, {ok, FFNIntermediate} = linear(CrossAttnNormed, W1, undefined), {ok, FFNActivated} = mlx:relu(FFNIntermediate), {ok, FFNOutput} = linear(FFNActivated, W2, undefined), {ok, FFNDropped} = dropout(FFNOutput, Dropout), {ok, FFNResidual} = mlx:add(CrossAttnNormed, FFNDropped), layer_norm(FFNResidual, LayerNorm3). positional_encoding(MaxLen, DModel, Base) -> positional_encoding(MaxLen, DModel, Base, float32). positional_encoding(MaxLen, DModel, Base, Dtype) -> % Sinusoidal positional encoding {ok, Position} = mlx:arange(0, MaxLen, 1), {ok, PosReshaped} = mlx:reshape(Position, [MaxLen, 1]), {ok, DimIndices} = mlx:arange(0, DModel div 2, 1), {ok, DivTerm} = mlx:power(mlx:array(Base), mlx:divide(mlx:multiply(DimIndices, mlx:array(2.0)), mlx:array(DModel))), {ok, Args} = mlx:divide(PosReshaped, DivTerm), {ok, SinEncoding} = mlx:sin(Args), {ok, CosEncoding} = mlx:cos(Args), % Interleave sin and cos {ok, Encoding} = mlx:stack([SinEncoding, CosEncoding], -1), {ok, EncodingFlat} = mlx:reshape(Encoding, [MaxLen, DModel]), case Dtype of float32 -> {ok, EncodingFlat}; _ -> mlx:cast(EncodingFlat, Dtype) end. %% Normalization layers layer_norm(Input, Weight) -> layer_norm(Input, Weight, undefined, 1.0e-5). layer_norm(Input, Weight, Bias) -> layer_norm(Input, Weight, Bias, 1.0e-5). layer_norm(Input, Weight, Bias, Eps) -> % Layer normalization {ok, Mean} = mlx:mean(Input, -1, true), {ok, Variance} = mlx:var(Input, -1, true), {ok, EpsArray} = mlx:array(Eps), {ok, VarPlusEps} = mlx:add(Variance, EpsArray), {ok, Std} = mlx:sqrt(VarPlusEps), {ok, Normalized} = mlx:divide(mlx:subtract(Input, Mean), Std), {ok, Scaled} = mlx:multiply(Normalized, Weight), case Bias of undefined -> {ok, Scaled}; _ -> mlx:add(Scaled, Bias) end. batch_norm(Input, Weight) -> batch_norm(Input, Weight, undefined, 1.0e-5). batch_norm(Input, Weight, Bias) -> batch_norm(Input, Weight, Bias, 1.0e-5). batch_norm(Input, Weight, Bias, Eps) -> % Batch normalization {ok, Mean} = mlx:mean(Input, 0, true), {ok, Variance} = mlx:var(Input, 0, true), {ok, EpsArray} = mlx:array(Eps), {ok, VarPlusEps} = mlx:add(Variance, EpsArray), {ok, Std} = mlx:sqrt(VarPlusEps), {ok, Normalized} = mlx:divide(mlx:subtract(Input, Mean), Std), {ok, Scaled} = mlx:multiply(Normalized, Weight), case Bias of undefined -> {ok, Scaled}; _ -> mlx:add(Scaled, Bias) end. group_norm(Input, NumGroups, Weight) -> group_norm(Input, NumGroups, Weight, undefined). group_norm(Input, NumGroups, Weight, Bias) -> % Group normalization mlx_nif:group_norm(Input, NumGroups, Weight, Bias). instance_norm(Input, Weight) -> instance_norm(Input, Weight, undefined). instance_norm(Input, Weight, Bias) -> % Instance normalization mlx_nif:instance_norm(Input, Weight, Bias). rms_norm(Input, Weight) -> rms_norm(Input, Weight, 1.0e-8). rms_norm(Input, Weight, Eps) -> % Root Mean Square normalization {ok, SquaredInput} = mlx:square(Input), {ok, MeanSquared} = mlx:mean(SquaredInput, -1, true), {ok, EpsArray} = mlx:array(Eps), {ok, MeanSquaredPlusEps} = mlx:add(MeanSquared, EpsArray), {ok, Rms} = mlx:sqrt(MeanSquaredPlusEps), {ok, Normalized} = mlx:divide(Input, Rms), mlx:multiply(Normalized, Weight). %% Dropout and regularization dropout(Input, Rate) -> dropout(Input, Rate, true). dropout(Input, Rate, Training) -> case Training andalso Rate > 0.0 of true -> {ok, Shape} = mlx:shape(Input), {ok, Mask} = random_bernoulli(Shape, 1.0 - Rate), {ok, Scale} = mlx:array(1.0 / (1.0 - Rate)), {ok, Dropped} = mlx:multiply(Input, Mask), mlx:multiply(Dropped, Scale); false -> {ok, Input} end. dropout2d(Input, Rate) -> mlx_nif:dropout2d(Input, Rate). dropout3d(Input, Rate) -> mlx_nif:dropout3d(Input, Rate). alpha_dropout(Input, Rate) -> mlx_nif:alpha_dropout(Input, Rate). feature_alpha_dropout(Input, Rate) -> mlx_nif:feature_alpha_dropout(Input, Rate). %% Advanced architectures residual_block(Input, Layers, Activation) -> residual_block(Input, Layers, Activation, #{}). residual_block(Input, Layers, Activation, Options) -> % Residual connection block Output = lists:foldl(fun(Layer, Acc) -> case Layer(Acc) of {ok, Result} -> Result; Result -> Result end end, Input, Layers), % Apply activation before residual connection {ok, Activated} = case Activation of relu -> mlx:relu(Output); gelu -> mlx_advanced:gelu(Output); swish -> mlx_advanced:swish(Output); identity -> {ok, Output} end, mlx:add(Input, Activated). dense_block(Input, Layers, GrowthRate) -> dense_block(Input, Layers, GrowthRate, #{}). dense_block(Input, Layers, GrowthRate, Options) -> % Dense connection block (DenseNet style) lists:foldl(fun(Layer, {Acc, Features}) -> {ok, NewFeature} = Layer(Acc), {ok, ConcatFeatures} = mlx:concatenate([Features, NewFeature], -1), {ConcatFeatures, ConcatFeatures} end, {Input, Input}, Layers). squeeze_excitation(Input, ReductionRatio) -> squeeze_excitation(Input, ReductionRatio, sigmoid). squeeze_excitation(Input, ReductionRatio, Activation) -> % Squeeze-and-Excitation block {ok, [B, C, H, W]} = mlx:shape(Input), % Squeeze: Global average pooling {ok, Squeezed} = global_avg_pool(Input), {ok, SqueezedFlat} = mlx:reshape(Squeezed, [B, C]), % Excitation: FC layers ReducedChannels = erlang:max(1, C div ReductionRatio), {ok, FC1Weight} = xavier_uniform([C, ReducedChannels]), {ok, FC2Weight} = xavier_uniform([ReducedChannels, C]), {ok, Excited1} = linear(SqueezedFlat, FC1Weight, undefined), {ok, Activated} = mlx:relu(Excited1), {ok, Excited2} = linear(Activated, FC2Weight, undefined), {ok, Scale} = case Activation of sigmoid -> mlx:sigmoid(Excited2); tanh -> mlx:tanh(Excited2) end, {ok, ScaleReshaped} = mlx:reshape(Scale, [B, C, 1, 1]), mlx:multiply(Input, ScaleReshaped). %% Recurrent layers lstm_cell(Input, Hidden, Cell, Weights) -> lstm_cell(Input, Hidden, Cell, Weights, #{}). lstm_cell(Input, Hidden, Cell, Weights, Options) -> % LSTM cell implementation [WIH, WHH, BIH, BHH] = Weights, % Compute gates {ok, GateInput} = linear(Input, WIH, BIH), {ok, GateHidden} = linear(Hidden, WHH, BHH), {ok, Gates} = mlx:add(GateInput, GateHidden), % Split gates {ok, [InputGate, ForgetGate, CellGate, OutputGate]} = mlx:split(Gates, 4, -1), % Apply activations {ok, InputActivated} = mlx:sigmoid(InputGate), {ok, ForgetActivated} = mlx:sigmoid(ForgetGate), {ok, CellActivated} = mlx:tanh(CellGate), {ok, OutputActivated} = mlx:sigmoid(OutputGate), % Update cell state {ok, ForgetComponent} = mlx:multiply(ForgetActivated, Cell), {ok, InputComponent} = mlx:multiply(InputActivated, CellActivated), {ok, NewCell} = mlx:add(ForgetComponent, InputComponent), % Compute new hidden state {ok, CellTanh} = mlx:tanh(NewCell), {ok, NewHidden} = mlx:multiply(OutputActivated, CellTanh), {ok, {NewHidden, NewCell}}. gru_cell(Input, Hidden, Weights) -> gru_cell(Input, Hidden, Weights, #{}). gru_cell(Input, Hidden, Weights, Options) -> % GRU cell implementation [WIH, WHH, BIH, BHH] = Weights, % Reset and update gates {ok, ResetUpdateInput} = linear(Input, WIH, BIH), {ok, ResetUpdateHidden} = linear(Hidden, WHH, BHH), {ok, ResetUpdate} = mlx:add(ResetUpdateInput, ResetUpdateHidden), {ok, [ResetGate, UpdateGate]} = mlx:split(ResetUpdate, 2, -1), {ok, ResetActivated} = mlx:sigmoid(ResetGate), {ok, UpdateActivated} = mlx:sigmoid(UpdateGate), % New gate {ok, ResetHidden} = mlx:multiply(ResetActivated, Hidden), {ok, NewInput} = linear(Input, WIH, BIH), {ok, NewHidden} = linear(ResetHidden, WHH, BHH), {ok, NewGate} = mlx:add(NewInput, NewHidden), {ok, NewActivated} = mlx:tanh(NewGate), % Final hidden state {ok, OneMinusUpdate} = mlx:subtract(mlx:array(1.0), UpdateActivated), {ok, HiddenComponent} = mlx:multiply(OneMinusUpdate, Hidden), {ok, NewComponent} = mlx:multiply(UpdateActivated, NewActivated), {ok, FinalHidden} = mlx:add(HiddenComponent, NewComponent), {ok, FinalHidden}. rnn_cell(Input, Hidden, Weights) -> rnn_cell(Input, Hidden, Weights, #{}). rnn_cell(Input, Hidden, Weights, Options) -> % Simple RNN cell [WIH, WHH, BIH, BHH] = Weights, {ok, InputComponent} = linear(Input, WIH, BIH), {ok, HiddenComponent} = linear(Hidden, WHH, BHH), {ok, Combined} = mlx:add(InputComponent, HiddenComponent), Activation = maps:get(activation, Options, tanh), case Activation of tanh -> mlx:tanh(Combined); relu -> mlx:relu(Combined); sigmoid -> mlx:sigmoid(Combined) end. %% Modern architectures vision_transformer_block(Input, Attention, MLP, LayerNorm1) -> vision_transformer_block(Input, Attention, MLP, LayerNorm1, undefined). vision_transformer_block(Input, Attention, MLP, LayerNorm1, LayerNorm2) -> % Vision Transformer block {ok, Normed1} = layer_norm(Input, LayerNorm1), {ok, AttnOutput} = Attention(Normed1), {ok, Residual1} = mlx:add(Input, AttnOutput), LayerNorm2Final = case LayerNorm2 of undefined -> LayerNorm1; _ -> LayerNorm2 end, {ok, Normed2} = layer_norm(Residual1, LayerNorm2Final), {ok, MLPOutput} = MLP(Normed2), mlx:add(Residual1, MLPOutput). swin_transformer_block(Input, WindowSize, Attention, MLP, LayerNorm1) -> swin_transformer_block(Input, WindowSize, Attention, MLP, LayerNorm1, undefined). swin_transformer_block(Input, WindowSize, Attention, MLP, LayerNorm1, LayerNorm2) -> % Swin Transformer block with windowed attention {ok, Windows} = window_partition(Input, WindowSize), {ok, AttnWindows} = Attention(Windows), {ok, Merged} = window_reverse(AttnWindows, WindowSize), vision_transformer_block(Input, fun(_) -> {ok, Merged} end, MLP, LayerNorm1, LayerNorm2). efficientnet_block(Input, ExpandRatio, KernelSize, Stride, SEReduction) -> efficientnet_block(Input, ExpandRatio, KernelSize, Stride, SEReduction, #{}). efficientnet_block(Input, ExpandRatio, KernelSize, Stride, SEReduction, Options) -> % EfficientNet MBConv block {ok, [_B, C, _H, _W]} = mlx:shape(Input), ExpandedChannels = C * ExpandRatio, % Expansion phase {ok, Expanded} = case ExpandRatio > 1 of true -> {ok, ExpWeight} = xavier_uniform([C, ExpandedChannels]), conv2d(Input, ExpWeight, undefined, 1); false -> {ok, Input} end, % Depthwise convolution {ok, DWWeight} = xavier_uniform([ExpandedChannels, 1, KernelSize, KernelSize]), {ok, DepthwiseOutput} = conv2d(Expanded, DWWeight, undefined, Stride), % Squeeze-and-Excitation {ok, SEOutput} = squeeze_excitation(DepthwiseOutput, SEReduction), % Projection phase {ok, ProjWeight} = xavier_uniform([ExpandedChannels, C]), {ok, ProjectionOutput} = conv2d(SEOutput, ProjWeight, undefined, 1), % Residual connection (if stride = 1 and same channels) case Stride =:= 1 andalso element(2, mlx:shape(Input)) =:= C of true -> mlx:add(Input, ProjectionOutput); false -> {ok, ProjectionOutput} end. %% Initialization functions xavier_uniform(Shape) -> % Xavier/Glorot uniform initialization [Fan_in, Fan_out] = case length(Shape) of 2 -> Shape; _ -> [lists:nth(1, Shape), lists:nth(2, Shape)] end, Limit = math:sqrt(6.0 / (Fan_in + Fan_out)), {ok, Random} = random_uniform(Shape, -Limit, Limit), {ok, Random}. xavier_normal(Shape) -> % Xavier/Glorot normal initialization [Fan_in, Fan_out] = case length(Shape) of 2 -> Shape; _ -> [lists:nth(1, Shape), lists:nth(2, Shape)] end, Std = math:sqrt(2.0 / (Fan_in + Fan_out)), random_normal(Shape, 0.0, Std). kaiming_uniform(Shape) -> % He/Kaiming uniform initialization Fan_in = lists:nth(1, Shape), Limit = math:sqrt(6.0 / Fan_in), random_uniform(Shape, -Limit, Limit). kaiming_normal(Shape) -> % He/Kaiming normal initialization Fan_in = lists:nth(1, Shape), Std = math:sqrt(2.0 / Fan_in), random_normal(Shape, 0.0, Std). orthogonal_init(Shape) -> % Orthogonal initialization mlx_nif:orthogonal_init(Shape). identity_init(Shape) -> % Identity initialization mlx:eye(lists:nth(1, Shape)). %% Helper functions random_uniform(Shape, Low, High) -> mlx_nif:random_uniform(Shape, Low, High). random_normal(Shape, Mean, Std) -> mlx_nif:random_normal(Shape, Mean, Std). random_bernoulli(Shape, Prob) -> mlx_nif:random_bernoulli(Shape, Prob). window_partition(Input, WindowSize) -> mlx_nif:window_partition(Input, WindowSize). window_reverse(Windows, WindowSize) -> mlx_nif:window_reverse(Windows, WindowSize).