diff --git a/docs/_static/braincore.jpg b/docs/_static/braincore.jpg deleted file mode 100644 index 6eca184..0000000 Binary files a/docs/_static/braincore.jpg and /dev/null differ diff --git a/docs/api.rst b/docs/api.rst index 84d4b8f..4fb71a9 100644 --- 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 new file mode 100644 index 0000000..f8fd9c9 --- /dev/null +++ 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 new file mode 100644 index 0000000..72323e1 --- /dev/null +++ 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 b/docs/apis/brainstate.rst similarity index 92% rename from docs/apis/braincore.rst rename to docs/apis/brainstate.rst index e027fe6..f990bc2 100644 --- a/docs/apis/braincore.rst +++ b/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 new file mode 100644 index 0000000..cf4398f --- /dev/null +++ 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 index 35ca35c..eaa654b 100644 --- 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 index 1be63a2..d8b30a6 100644 --- 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/ diff --git a/docs/apis/optim.rst b/docs/apis/optim.rst new file mode 100644 index 0000000..191d1f2 --- /dev/null +++ b/docs/apis/optim.rst @@ -0,0 +1,33 @@ +``brainstate.optim`` module +=========================== + +.. currentmodule:: brainstate.optim +.. automodule:: brainstate.optim + +.. autosummary:: + :toctree: generated/ + + to_same_dict_tree + LearningRateScheduler + ConstantLR + StepLR + MultiStepLR + CosineAnnealingLR + CosineAnnealingWarmRestarts + ExponentialLR + ExponentialDecayLR + InverseTimeDecayLR + PolynomialDecayLR + PiecewiseConstantLR + OptimState + Optimizer + SGD + Momentum + MomentumNesterov + Adagrad + Adadelta + RMSProp + Adam + LARS + Adan + AdamW diff --git a/docs/apis/random.rst b/docs/apis/random.rst index 8b26b8d..ee99c67 100644 --- a/docs/apis/random.rst +++ b/docs/apis/random.rst @@ -1,8 +1,8 @@ -``braincore.random`` module -=========================== +``brainstate.random`` module +============================ -.. currentmodule:: braincore.random -.. automodule:: braincore.random +.. currentmodule:: brainstate.random +.. automodule:: brainstate.random Random Number Generators diff --git a/docs/apis/surrogate.rst b/docs/apis/surrogate.rst index 0148b78..d1255b4 100644 --- a/docs/apis/surrogate.rst +++ b/docs/apis/surrogate.rst @@ -1,8 +1,8 @@ -``braincore.surrogate`` module -============================== +``brainstate.surrogate`` module +=============================== -.. currentmodule:: braincore.surrogate -.. automodule:: braincore.surrogate +.. currentmodule:: brainstate.surrogate +.. automodule:: brainstate.surrogate Surrogate Gradient Functions diff --git a/docs/apis/transform.rst b/docs/apis/transform.rst index 045ef23..7e8fe10 100644 --- a/docs/apis/transform.rst +++ b/docs/apis/transform.rst @@ -1,8 +1,8 @@ -``braincore.transform`` module -============================== +``brainstate.transform`` module +=============================== -.. currentmodule:: braincore.transform -.. automodule:: braincore.transform +.. currentmodule:: brainstate.transform +.. automodule:: brainstate.transform diff --git a/docs/apis/util.rst b/docs/apis/util.rst index c9dfa2b..0014912 100644 --- a/docs/apis/util.rst +++ b/docs/apis/util.rst @@ -1,8 +1,8 @@ -``braincore.util`` module -========================= +``brainstate.util`` module +========================== -.. currentmodule:: braincore.util -.. automodule:: braincore.util +.. currentmodule:: brainstate.util +.. automodule:: brainstate.util .. autosummary:: :toctree: generated/ diff --git a/docs/auto_generater.py b/docs/auto_generater.py index 5d19805..0dface5 100644 --- a/docs/auto_generater.py +++ b/docs/auto_generater.py @@ -449,38 +449,75 @@ def generate_algorithm_docs(path='apis/auto/algorithms/'): def main(): os.makedirs('apis/auto/', exist_ok=True) - _write_module(module_name='brainstate.surrogate', - filename='apis/auto/surrogate.rst', - header='``brainstate.surrogate`` module') + # _write_module(module_name='brainstate.surrogate', + # filename='apis/auto/surrogate.rst', + # header='``brainstate.surrogate`` module') + + # _write_module(module_name='brainstate.random', + # filename='apis/auto/random.rst', + # header='``brainstate.random`` module') + + # _write_module(module_name='brainstate.mixin', + # filename='apis/auto/mixin.rst', + # header='``brainstate.mixin`` module') + + # _write_module(module_name='brainstate.transform', + # filename='apis/auto/transform.rst', + # header='``brainstate.transform`` module') + + # _write_module(module_name='brainstate.math', + # filename='apis/auto/math.rst', + # header='``brainstate.math`` module') + + # _write_module(module_name='brainstate.util', + # filename='apis/auto/util.rst', + # header='``brainstate.util`` module') + + _write_module(module_name='brainstate.optim', + filename='apis/optim.rst', + header='``brainstate.optim`` module') + + _write_module(module_name='brainstate.init', + filename='apis/init.rst', + header='``brainstate.init`` module') + + # module_and_name = [ + # ('_state', 'State System'), + # ('_module', 'Module & Container'), + # ] + # _write_submodules(module_name='brainstate', + # filename='apis/auto/brainstate.rst', + # header='``brainstate`` module', + # submodule_names=[k[0] for k in module_and_name], + # section_names=[k[1] for k in module_and_name]) - _write_module(module_name='brainstate.random', - filename='apis/auto/random.rst', - header='``brainstate.random`` module') - - _write_module(module_name='brainstate.mixin', - filename='apis/auto/mixin.rst', - header='``brainstate.mixin`` module') - - _write_module(module_name='brainstate.transform', - filename='apis/auto/transform.rst', - header='``brainstate.transform`` module') - - _write_module(module_name='brainstate.math', - filename='apis/auto/math.rst', - header='``brainstate.math`` module') - - _write_module(module_name='brainstate.util', - filename='apis/auto/util.rst', - header='``brainstate.util`` module') + module_and_name = [ + ('_activations', 'Activation Functions'), + ('_normalization', 'Normalization Functions'), + ('_spikes', 'Spiking Operations'), + ] + _write_submodules(module_name='brainstate.functional', + filename='apis/brainstate.functional.rst', + header='``brainstate.functional`` module', + submodule_names=[k[0] for k in module_and_name], + section_names=[k[1] for k in module_and_name]) module_and_name = [ - ('_state', 'State System'), - ('_module', 'Module & Container'), + ('_base', 'Base Classes'), ('_projection', 'Synaptic Projections'), + ('_connections', 'Connection Layers'), + ('_dynamics', 'Neuronal/Synaptic Dynamics'), + ('_rate_rnns', 'Rate RNNs'), + ('_readout', 'Readout Layers'), + ('_synouts', 'Synaptic Outputs'), + ('_elementwise', 'Element-wise Layers'), + ('_normalizations', 'Normalization Layers'), + ('_poolings', 'Pooling Layers'), + ('_others', 'Other Layers'), ] - _write_submodules(module_name='brainstate', - filename='apis/auto/brainstate.rst', - header='``brainstate`` module', + _write_submodules(module_name='brainstate.nn', + filename='apis/brainstate.nn.rst', + header='``brainstate.nn`` module', submodule_names=[k[0] for k in module_and_name], section_names=[k[1] for k in module_and_name])