Skip to content

Commit

Permalink
merge general linear to linear, remove FNN
Browse files Browse the repository at this point in the history
  • Loading branch information
ASEM000 committed Nov 28, 2023
1 parent aedf891 commit e829d93
Show file tree
Hide file tree
Showing 6 changed files with 60 additions and 431 deletions.
3 changes: 0 additions & 3 deletions docs/API/linear.rst
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,5 @@ Linear

.. autoclass:: Linear
.. autoclass:: Identity
.. autoclass:: GeneralLinear
.. autoclass:: Embedding

.. autoclass:: FNN
.. autoclass:: MLP
3 changes: 0 additions & 3 deletions serket/_src/nn/activation.py
Original file line number Diff line number Diff line change
Expand Up @@ -561,9 +561,6 @@ def def_act_entry(key: str, act: ActivationFunctionType) -> None:
... def __call__(self, x):
... return x * self.my_param
>>> sk.def_act_entry("my_act", MyTrainableActivation())
>>> x = jnp.ones((1, 1))
>>> sk.nn.FNN([1, 1, 1], act="my_act", weight_init="ones", bias_init=None, key=jr.PRNGKey(0))(x)
Array([[10.]], dtype=float32)
"""
if key in act_map:
raise ValueError(f"`init_key` {key=} already registered")
Expand Down
Loading

0 comments on commit e829d93

Please sign in to comment.