update documentation

chaobrain · Jun 9, 2024 · 75c0b3d · 75c0b3d
1 parent 1f7c789
commit 75c0b3d
Show file tree

Hide file tree

Showing 14 changed files with 454 additions and 60 deletions.
diff --git a/docs/_static/braincore.jpg b/docs/_static/braincore.jpg
diff --git a/docs/api.rst b/docs/api.rst
@@ -5,10 +5,14 @@ API Documentation
    :maxdepth: 1
 
    apis/changelog.md
-   apis/braincore.rst
+   apis/brainstate.rst
+   apis/init.rst
    apis/math.rst
    apis/mixin.rst
+   apis/optim.rst
    apis/random.rst
    apis/surrogate.rst
    apis/transform.rst
    apis/util.rst
+   apis/brainstate.nn.rst
+   apis/brainstate.functional.rst
diff --git a/docs/apis/brainstate.functional.rst b/docs/apis/brainstate.functional.rst
@@ -0,0 +1,76 @@
+``brainstate.functional`` module
+================================
+
+.. currentmodule:: brainstate.functional 
+.. automodule:: brainstate.functional 
+
+Activation Functions
+--------------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   tanh
+   relu
+   squareplus
+   softplus
+   soft_sign
+   sigmoid
+   silu
+   swish
+   log_sigmoid
+   elu
+   leaky_relu
+   hard_tanh
+   celu
+   selu
+   gelu
+   glu
+   logsumexp
+   log_softmax
+   softmax
+   standardize
+   one_hot
+   relu6
+   hard_sigmoid
+   hard_silu
+   hard_swish
+   hard_shrink
+   rrelu
+   mish
+   soft_shrink
+   prelu
+   tanh_shrink
+   softmin
+
+
+Normalization Functions
+-----------------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   weight_standardization
+
+
+Spiking Operations
+------------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   spike_bitwise_or
+   spike_bitwise_and
+   spike_bitwise_iand
+   spike_bitwise_not
+   spike_bitwise_xor
+   spike_bitwise_ixor
+   spike_bitwise
+
+
diff --git a/docs/apis/brainstate.nn.rst b/docs/apis/brainstate.nn.rst
@@ -0,0 +1,215 @@
+``brainstate.nn`` module
+========================
+
+.. currentmodule:: brainstate.nn 
+.. automodule:: brainstate.nn 
+
+Base Classes
+------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   ExplicitInOutSize
+   ElementWiseBlock
+   Sequential
+   DnnLayer
+
+
+Synaptic Projections
+--------------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   HalfProjAlignPostMg
+   FullProjAlignPostMg
+   HalfProjAlignPost
+   FullProjAlignPost
+   FullProjAlignPreSDMg
+   FullProjAlignPreDSMg
+   FullProjAlignPreSD
+   FullProjAlignPreDS
+   HalfProjDelta
+   FullProjDelta
+   VanillaProj
+
+
+Connection Layers
+-----------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   Linear
+   ScaledWSLinear
+   SignedWLinear
+   CSRLinear
+   Conv1d
+   Conv2d
+   Conv3d
+   ScaledWSConv1d
+   ScaledWSConv2d
+   ScaledWSConv3d
+
+
+Neuronal/Synaptic Dynamics
+--------------------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   Neuron
+   IF
+   LIF
+   ALIF
+   Synapse
+   Expon
+   STP
+   STD
+
+
+Rate RNNs
+---------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   RNNCell
+   ValinaRNNCell
+   GRUCell
+   MGUCell
+   LSTMCell
+   URLSTMCell
+
+
+Readout Layers
+--------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   LeakyRateReadout
+   LeakySpikeReadout
+
+
+Synaptic Outputs
+----------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   SynOut
+   COBA
+   CUBA
+   MgBlock
+
+
+Element-wise Layers
+-------------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   Threshold
+   ReLU
+   RReLU
+   Hardtanh
+   ReLU6
+   Sigmoid
+   Hardsigmoid
+   Tanh
+   SiLU
+   Mish
+   Hardswish
+   ELU
+   CELU
+   SELU
+   GLU
+   GELU
+   Hardshrink
+   LeakyReLU
+   LogSigmoid
+   Softplus
+   Softshrink
+   PReLU
+   Softsign
+   Tanhshrink
+   Softmin
+   Softmax
+   Softmax2d
+   LogSoftmax
+   Dropout
+   Dropout1d
+   Dropout2d
+   Dropout3d
+   AlphaDropout
+   FeatureAlphaDropout
+   Identity
+   SpikeBitwise
+
+
+Normalization Layers
+--------------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   BatchNorm1d
+   BatchNorm2d
+   BatchNorm3d
+
+
+Pooling Layers
+--------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   Flatten
+   Unflatten
+   AvgPool1d
+   AvgPool2d
+   AvgPool3d
+   MaxPool1d
+   MaxPool2d
+   MaxPool3d
+   AdaptiveAvgPool1d
+   AdaptiveAvgPool2d
+   AdaptiveAvgPool3d
+   AdaptiveMaxPool1d
+   AdaptiveMaxPool2d
+   AdaptiveMaxPool3d
+
+
+Other Layers
+------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+   :template: classtemplate.rst
+
+   DropoutFixed
+
+
diff --git a/docs/apis/braincore.rst → docs/apis/brainstate.rst b/docs/apis/braincore.rst → docs/apis/brainstate.rst
@@ -1,8 +1,8 @@
-``braincore`` module
-====================
+``brainstate`` module
+=====================
 
-.. currentmodule:: braincore 
-.. automodule:: braincore 
+.. currentmodule:: brainstate 
+.. automodule:: brainstate 
 
 ``State`` System
 ----------------

diff --git a/docs/apis/init.rst b/docs/apis/init.rst
@@ -0,0 +1,29 @@
+``brainstate.init`` module
+==========================
+
+.. currentmodule:: brainstate.init 
+.. automodule:: brainstate.init 
+
+.. autosummary::
+   :toctree: generated/
+
+   param
+   state
+   noise
+   to_size
+   Initializer
+   ZeroInit
+   Constant
+   Identity
+   Normal
+   TruncatedNormal
+   Uniform
+   VarianceScaling
+   KaimingUniform
+   KaimingNormal
+   XavierUniform
+   XavierNormal
+   LecunUniform
+   LecunNormal
+   Orthogonal
+   DeltaOrthogonal
diff --git a/docs/apis/math.rst b/docs/apis/math.rst
@@ -1,8 +1,8 @@
-``braincore.math`` module
-=========================
+``brainstate.math`` module
+==========================
 
-.. currentmodule:: braincore.math 
-.. automodule:: braincore.math 
+.. currentmodule:: brainstate.math 
+.. automodule:: brainstate.math 
 
 .. autosummary::
    :toctree: generated/
@@ -19,7 +19,7 @@
    as_numpy
    tree_zeros_like
    tree_ones_like
-   ein_reduce
-   ein_rearrange
-   ein_repeat
-   ein_shape
+   einreduce
+   einrearrange
+   einrepeat
+   einshape
diff --git a/docs/apis/mixin.rst b/docs/apis/mixin.rst
@@ -1,8 +1,8 @@
-``braincore.mixin`` module
-==========================
+``brainstate.mixin`` module
+===========================
 
-.. currentmodule:: braincore.mixin 
-.. automodule:: braincore.mixin 
+.. currentmodule:: brainstate.mixin 
+.. automodule:: brainstate.mixin 
 
 .. autosummary::
    :toctree: generated/