move to _src

docs/notebooks/layers_overview.ipynb

@@ -4,7 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 📙 `serket.nn` layers overview"
+    "# 📙 `serket._src.nn` layers overview"
    ]
   },
   {

serket/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -42,10 +42,11 @@
     unfreeze,
 )
 
+from serket._src.custom_transform import tree_eval, tree_state
+from serket._src.nn.activation import def_act_entry
+from serket._src.nn.initialization import def_init_entry
+
 from . import cluster, image, nn
-from .custom_transform import tree_eval, tree_state
-from .nn.activation import def_act_entry
-from .nn.initialization import def_init_entry
 
 __all__ = [
     # general utils

serket/_src/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2023 serket authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

serket/_src/cluster/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2023 serket authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

serket/cluster/kmeans.py → serket/_src/cluster/kmeans.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -23,8 +23,8 @@
 from typing_extensions import Annotated
 
 import serket as sk
-from serket.custom_transform import tree_eval, tree_state
-from serket.utils import IsInstance, Range
+from serket._src.custom_transform import tree_eval, tree_state
+from serket._src.utils import IsInstance, Range
 
 """K-means utility functions."""
 

serket/custom_transform.py → serket/_src/custom_transform.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -133,7 +133,7 @@ def tree_eval(tree):
 
     Note:
         To define evaluation rule for a custom layer, use the decorator
-        :func:`.tree_eval.def_eval` on a function that accepts the layer. The 
+        :func:`.tree_eval.def_eval` on a function that accepts the layer. The
         function should return the evaluation layer.
 
         >>> import serket as sk

serket/_src/image/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2023 serket authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

serket/image/augment.py → serket/_src/image/augment.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -22,9 +22,9 @@
 from jax import lax
 
 import serket as sk
-from serket.custom_transform import tree_eval
-from serket.nn.linear import Identity
-from serket.utils import IsInstance, Range, validate_spatial_ndim
+from serket._src.custom_transform import tree_eval
+from serket._src.nn.linear import Identity
+from serket._src.utils import IsInstance, Range, validate_spatial_ndim
 
 
 def pixel_shuffle_2d(x: jax.Array, upscale_factor: int | tuple[int, int]) -> jax.Array:

serket/image/filter.py → serket/_src/image/filter.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -21,9 +21,9 @@
 from typing_extensions import Annotated
 
 import serket as sk
-from serket.nn.convolution import fft_conv_general_dilated
-from serket.nn.initialization import DType
-from serket.utils import (
+from serket._src.nn.convolution import fft_conv_general_dilated
+from serket._src.nn.initialization import DType
+from serket._src.utils import (
     generate_conv_dim_numbers,
     positive_int_cb,
     resolve_string_padding,

serket/image/geometric.py → serket/_src/image/geometric.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -22,9 +22,9 @@
 from jax.scipy.ndimage import map_coordinates
 
 import serket as sk
-from serket.custom_transform import tree_eval
-from serket.nn.linear import Identity
-from serket.utils import IsInstance, validate_spatial_ndim
+from serket._src.custom_transform import tree_eval
+from serket._src.nn.linear import Identity
+from serket._src.utils import IsInstance, validate_spatial_ndim
 
 
 def affine(image, matrix):

serket/_src/nn/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2023 serket authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

serket/nn/activation.py → serket/_src/nn/activation.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -22,7 +22,7 @@
 from jax import lax
 
 import serket as sk
-from serket.utils import IsInstance, Range, ScalarLike
+from serket._src.utils import IsInstance, Range, ScalarLike
 
 T = TypeVar("T")
 

serket/nn/attention.py → serket/_src/nn/attention.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -22,8 +22,8 @@
 from typing_extensions import Annotated
 
 import serket as sk
-from serket.nn.initialization import InitType
-from serket.utils import maybe_lazy_call, maybe_lazy_init
+from serket._src.nn.initialization import InitType
+from serket._src.utils import maybe_lazy_call, maybe_lazy_init
 
 """Defines attention layers."""
 

serket/nn/containers.py → serket/_src/nn/containers.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -21,8 +21,8 @@
 import jax.random as jr
 
 import serket as sk
-from serket.custom_transform import tree_eval
-from serket.utils import Range
+from serket._src.custom_transform import tree_eval
+from serket._src.utils import Range
 
 
 def sequential(

serket/nn/convolution.py → serket/_src/nn/convolution.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -27,8 +27,8 @@
 from typing_extensions import Annotated
 
 import serket as sk
-from serket.nn.initialization import DType, InitType, resolve_init
-from serket.utils import (
+from serket._src.nn.initialization import DType, InitType, resolve_init
+from serket._src.utils import (
     DilationType,
     KernelSizeType,
     PaddingType,

serket/nn/dropout.py → serket/_src/nn/dropout.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -23,9 +23,8 @@
 import jax.random as jr
 
 import serket as sk
-from serket.custom_transform import tree_eval
-from serket.nn.linear import Identity
-from serket.utils import (
+from serket._src.custom_transform import tree_eval
+from serket._src.utils import (
     IsInstance,
     Range,
     canonicalize,
@@ -434,5 +433,5 @@ def spatial_ndim(self) -> int:
 @tree_eval.def_eval(RandomCutout2D)
 @tree_eval.def_eval(GeneralDropout)
 @tree_eval.def_eval(DropoutND)
-def dropout_evaluation(_) -> Identity:
-    return Identity()
+def dropout_evaluation(_) -> sk.nn.Identity:
+    return sk.nn.Identity()

serket/nn/initialization.py → serket/_src/nn/initialization.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,7 +15,7 @@
 
 import functools as ft
 from collections.abc import Callable as ABCCallable
-from typing import Callable,Any, Literal, Tuple, Union, get_args
+from typing import Any, Callable, Literal, Tuple, Union, get_args
 
 import jax
 import jax.nn.initializers as ji

serket/nn/linear.py → serket/_src/nn/linear.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -22,13 +22,13 @@
 import jax.random as jr
 
 import serket as sk
-from serket.nn.activation import (
+from serket._src.nn.activation import (
     ActivationFunctionType,
     ActivationType,
     resolve_activation,
 )
-from serket.nn.initialization import DType, InitType, resolve_init
-from serket.utils import maybe_lazy_call, maybe_lazy_init, positive_int_cb
+from serket._src.nn.initialization import DType, InitType, resolve_init
+from serket._src.utils import maybe_lazy_call, maybe_lazy_init, positive_int_cb
 
 T = TypeVar("T")
 

serket/nn/normalization.py → serket/_src/nn/normalization.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -22,9 +22,9 @@
 from jax.custom_batching import custom_vmap
 
 import serket as sk
-from serket.custom_transform import tree_eval, tree_state
-from serket.nn.initialization import DType, InitType, resolve_init
-from serket.utils import (
+from serket._src.custom_transform import tree_eval, tree_state
+from serket._src.nn.initialization import DType, InitType, resolve_init
+from serket._src.utils import (
     Range,
     ScalarLike,
     maybe_lazy_call,
@@ -225,6 +225,7 @@ class GroupNorm(sk.TreeClass):
     Reference:
         https://nn.labml.ai/normalization/group_norm/index.html
     """
+
     eps: float = sk.field(on_setattr=[Range(0), ScalarLike()])
 
     @ft.partial(maybe_lazy_init, is_lazy=is_lazy_init)

serket/nn/pooling.py → serket/_src/nn/pooling.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -23,7 +23,7 @@
 import kernex as kex
 
 import serket as sk
-from serket.utils import (
+from serket._src.utils import (
     KernelSizeType,
     PaddingType,
     StridesType,

serket/nn/recurrent.py → serket/_src/nn/recurrent.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -24,10 +24,10 @@
 import jax.tree_util as jtu
 
 import serket as sk
-from serket.custom_transform import tree_state
-from serket.nn.activation import ActivationType, resolve_activation
-from serket.nn.initialization import DType, InitType
-from serket.utils import (
+from serket._src.custom_transform import tree_state
+from serket._src.nn.activation import ActivationType, resolve_activation
+from serket._src.nn.initialization import DType, InitType
+from serket._src.utils import (
     DilationType,
     KernelSizeType,
     PaddingType,

serket/nn/reshape.py → serket/_src/nn/reshape.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -23,9 +23,9 @@
 import jax.random as jr
 
 import serket as sk
-from serket.custom_transform import tree_eval
-from serket.nn.linear import Identity
-from serket.utils import (
+from serket._src.custom_transform import tree_eval
+from serket._src.nn.linear import Identity
+from serket._src.utils import (
     IsInstance,
     canonicalize,
     delayed_canonicalize_padding,

serket/utils.py → serket/_src/utils.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Serket authors
+# Copyright 2023 serket authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.