__mo_san__
a123c3c7d2
- Add batch normalization layer --------- Co-authored-by: Robert McCraith <mccraithrobert@gmail.com> Co-authored-by: Awni Hannun <awni@apple.com>
34 lines
374 B
ReStructuredText
34 lines
374 B
ReStructuredText
.. _layers:
|
|
|
|
.. currentmodule:: mlx.nn
|
|
|
|
Layers
|
|
------
|
|
|
|
.. autosummary::
|
|
:toctree: _autosummary
|
|
:template: nn-module-template.rst
|
|
|
|
Embedding
|
|
ReLU
|
|
PReLU
|
|
GELU
|
|
SiLU
|
|
Step
|
|
SELU
|
|
Mish
|
|
Linear
|
|
Conv1d
|
|
Conv2d
|
|
BatchNorm
|
|
LayerNorm
|
|
RMSNorm
|
|
GroupNorm
|
|
RoPE
|
|
MultiHeadAttention
|
|
Sequential
|
|
QuantizedLinear
|
|
Dropout
|
|
Dropout2d
|
|
|