diff --git a/docs/assets/images/black_logo_description.png b/docs/assets/images/black_logo_description.png new file mode 100644 index 0000000..83e1425 Binary files /dev/null and b/docs/assets/images/black_logo_description.png differ diff --git a/docs/assets/images/black_logo_illia.png b/docs/assets/images/black_logo_illia.png index 28d30bb..5cc2ba3 100644 Binary files a/docs/assets/images/black_logo_illia.png and b/docs/assets/images/black_logo_illia.png differ diff --git a/docs/assets/images/logo_black_small.png b/docs/assets/images/logo_black_small.png deleted file mode 100644 index 0705c2d..0000000 Binary files a/docs/assets/images/logo_black_small.png and /dev/null differ diff --git a/docs/assets/images/white_logo_description.png b/docs/assets/images/white_logo_description.png new file mode 100644 index 0000000..661442c Binary files /dev/null and b/docs/assets/images/white_logo_description.png differ diff --git a/docs/assets/images/white_logo_illia.png b/docs/assets/images/white_logo_illia.png index f382068..695c08e 100644 Binary files a/docs/assets/images/white_logo_illia.png and b/docs/assets/images/white_logo_illia.png differ diff --git a/docs/examples/Computer Vision/MNIST Bayesian CNN.ipynb b/docs/examples/Computer Vision/MNIST Bayesian CNN.ipynb index 97104f3..dad1b02 100644 --- a/docs/examples/Computer Vision/MNIST Bayesian CNN.ipynb +++ b/docs/examples/Computer Vision/MNIST Bayesian CNN.ipynb @@ -18,7 +18,7 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "b3342740", "metadata": {}, "outputs": [ @@ -26,7 +26,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "0.2.1\n", + "0.2.4\n", "torch\n" ] } @@ -39,7 +39,7 @@ "os.environ[\"ILLIA_BACKEND\"] = \"torch\"\n", "\n", "import illia\n", - "from illia.nn import Conv2d, Linear\n", + "from illia.nn import Conv2d, Linear, ReLU, MaxPool2d, Dropout\n", "\n", "import numpy as np\n", "import torch\n", @@ -65,7 +65,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 13, "id": "1d0963f5", "metadata": {}, "outputs": [ @@ -92,7 +92,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 14, "id": "83ba2071", "metadata": {}, "outputs": [], @@ -125,7 +125,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "id": "c96c969c", "metadata": {}, "outputs": [ @@ -150,23 +150,27 @@ " self.fc1 = Linear(64 * 7 * 7, 128)\n", " self.fc2 = Linear(128, 10)\n", "\n", + " # Activation and pooling\n", + " self.relu = ReLU()\n", + " self.pool = MaxPool2d(2)\n", + "\n", " # Dropout for regularization\n", - " self.dropout = nn.Dropout(0.25)\n", + " self.dropout = Dropout(0.25)\n", "\n", " def forward(self, x):\n", " # First conv + ReLU + MaxPool\n", - " x = F.relu(self.conv1(x))\n", - " x = F.max_pool2d(x, 2)\n", + " x = self.relu(self.conv1(x))\n", + " x = self.pool(x)\n", "\n", " # Second conv + ReLU + MaxPool\n", - " x = F.relu(self.conv2(x))\n", - " x = F.max_pool2d(x, 2)\n", + " x = self.relu(self.conv2(x))\n", + " x = self.pool(x)\n", "\n", " # Flatten before fully connected layers\n", " x = x.view(x.size(0), -1)\n", "\n", " # Fully connected + dropout\n", - " x = F.relu(self.fc1(x))\n", + " x = self.relu(self.fc1(x))\n", " x = self.dropout(x)\n", " x = self.fc2(x)\n", "\n", @@ -189,7 +193,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 16, "id": "4b0b9fed", "metadata": {}, "outputs": [], @@ -212,7 +216,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 17, "id": "b80d4785", "metadata": {}, "outputs": [], @@ -285,7 +289,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, "id": "f7021581", "metadata": {}, "outputs": [ @@ -297,27 +301,9 @@ "\n", "Epoch 1/2\n", "--------------------------------------------------\n", - "Batch 0/938, Loss: 3.542201, Accuracy: 9.38%\n", - "Batch 200/938, Loss: 0.405245, Accuracy: 83.38%\n", - "Batch 400/938, Loss: 0.096597, Accuracy: 89.36%\n", - "Batch 600/938, Loss: 0.062778, Accuracy: 91.73%\n", - "Batch 800/938, Loss: 0.162707, Accuracy: 93.07%\n", - "\n", - "Test Loss: 0.043868, Test Accuracy: 98.52%\n", - "\n", - "Epoch 1 - Train Loss: 0.206314, Train Acc: 93.69%, Test Acc: 98.52%\n", - "Epoch 2/2\n", - "--------------------------------------------------\n", - "Batch 0/938, Loss: 0.052432, Accuracy: 98.44%\n", - "Batch 200/938, Loss: 0.013136, Accuracy: 98.00%\n", - "Batch 400/938, Loss: 0.008689, Accuracy: 98.06%\n", - "Batch 600/938, Loss: 0.035820, Accuracy: 98.03%\n", - "Batch 800/938, Loss: 0.024955, Accuracy: 98.05%\n", - "\n", - "Test Loss: 0.033828, Test Accuracy: 98.91%\n", - "\n", - "Epoch 2 - Train Loss: 0.060348, Train Acc: 98.10%, Test Acc: 98.91%\n", - "Training completed!\n" + "Batch 0/938, Loss: 5.569467, Accuracy: 12.50%\n", + "Batch 200/938, Loss: 0.297848, Accuracy: 82.00%\n", + "Batch 400/938, Loss: 0.080678, Accuracy: 88.43%\n" ] } ], @@ -365,7 +351,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": null, "id": "45bc2223", "metadata": {}, "outputs": [ @@ -487,7 +473,7 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": null, "id": "009a2b4d", "metadata": {}, "outputs": [], @@ -513,7 +499,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": null, "id": "e21d3a10", "metadata": {}, "outputs": [], @@ -553,7 +539,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": null, "id": "21e1f116", "metadata": {}, "outputs": [ @@ -596,7 +582,7 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": null, "id": "93e284d5", "metadata": {}, "outputs": [ @@ -642,7 +628,7 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": null, "id": "f5c224c5", "metadata": {}, "outputs": [ diff --git a/docs/index.md b/docs/index.md index 3d52e6e..158dbc0 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,6 +1,6 @@

- - + +

@@ -79,6 +79,15 @@ print(f"Output std: {outputs.std()}") print(f"Output var: {outputs.var()}") ``` +## Non-Parametric Definitions + +**illia** provides non-parametric layers (pooling, activation, normalization, regularization, and utility layers) imported from `illia.nn` using PyTorch-style naming conventions. However, the initialization parameters and API usage are backend-specific and not standardized across PyTorch, TensorFlow, and JAX. + +Key characteristics: +- **PyTorch Naming**: Layer names follow PyTorch conventions (e.g., MaxPool2d, ReLU, BatchNorm2d) +- **Backend-Specific Parameters**: Initialization arguments vary by backend +- **Non-Standard API**: JAX uses functional definitions, while PyTorch and TensorFlow use class/object-oriented layers + ## Contributing We welcome contributions from the community! Whether you're fixing bugs, adding diff --git a/illia/nn/__init__.py b/illia/nn/__init__.py index 39f14d9..2c396a6 100644 --- a/illia/nn/__init__.py +++ b/illia/nn/__init__.py @@ -9,29 +9,19 @@ from illia import BackendManager -# Obtain the library to import -def __getattr__(name: str) -> None: +def __getattr__(name: str) -> Any: """ - Dynamically import a class from backend distributions. + Dynamically import a class from backend. Args: name: Name of the class to be imported. Returns: - None. + The requested layer/module class. """ + backend = BackendManager.get_backend() + module = BackendManager.get_backend_module(backend, "nn") + layer_class = BackendManager.get_class(backend, name, "nn", module) - # Obtain parameters for nn - module_type: str = "nn" - backend: str = BackendManager.get_backend() - module_path: Any | dict[str, Any] = BackendManager.get_backend_module( - backend, module_type - ) - - # Set class to global namespace - globals()[name] = BackendManager.get_class( - backend_name=backend, - class_name=name, - module_type=module_type, - module_path=module_path, - ) + globals()[name] = layer_class + return layer_class diff --git a/illia/nn/jax/__init__.py b/illia/nn/jax/__init__.py index a650c06..f47d30b 100644 --- a/illia/nn/jax/__init__.py +++ b/illia/nn/jax/__init__.py @@ -5,19 +5,35 @@ """ # Own modules +from illia.nn.jax.activation import GELU, ReLU, Sigmoid, Tanh from illia.nn.jax.base import BayesianModule from illia.nn.jax.conv1d import Conv1d from illia.nn.jax.conv2d import Conv2d from illia.nn.jax.embedding import Embedding from illia.nn.jax.linear import Linear from illia.nn.jax.lstm import LSTM +from illia.nn.jax.normalization import BatchNorm1d, BatchNorm2d, LayerNorm +from illia.nn.jax.pooling import AvgPool1d, AvgPool2d, MaxPool1d, MaxPool2d +from illia.nn.jax.regularization import Dropout __all__: list[str] = [ + "AvgPool1d", + "AvgPool2d", + "BatchNorm1d", + "BatchNorm2d", "BayesianModule", "Conv1d", "Conv2d", + "Dropout", "Embedding", - "Linear", + "GELU", "LSTM", + "LayerNorm", + "Linear", + "MaxPool1d", + "MaxPool2d", + "ReLU", + "Sigmoid", + "Tanh", ] diff --git a/illia/nn/jax/activation.py b/illia/nn/jax/activation.py new file mode 100644 index 0000000..ded0975 --- /dev/null +++ b/illia/nn/jax/activation.py @@ -0,0 +1,35 @@ +"""JAX activation layer wrappers.""" + +# Standard libraries +from functools import partial + +# 3pps +from flax import nnx + + +class ReLU: + """Wrapper for JAX relu.""" + + def __new__(cls, *args, **kwargs): + return partial(nnx.relu, *args, **kwargs) + + +class Sigmoid: + """Wrapper for JAX sigmoid.""" + + def __new__(cls, *args, **kwargs): + return partial(nnx.sigmoid, *args, **kwargs) + + +class Tanh: + """Wrapper for JAX tanh.""" + + def __new__(cls, *args, **kwargs): + return partial(nnx.tanh, *args, **kwargs) + + +class GELU: + """Wrapper for JAX gelu.""" + + def __new__(cls, *args, **kwargs): + return partial(nnx.gelu, *args, **kwargs) diff --git a/illia/nn/jax/normalization.py b/illia/nn/jax/normalization.py new file mode 100644 index 0000000..58ca052 --- /dev/null +++ b/illia/nn/jax/normalization.py @@ -0,0 +1,25 @@ +"""JAX normalization layer wrappers.""" + +# 3pps +from flax import nnx + + +class BatchNorm1d: + """Wrapper for JAX BatchNorm for 1D inputs.""" + + def __new__(cls, num_features, *args, **kwargs): + return nnx.BatchNorm(num_features, *args, **kwargs) + + +class BatchNorm2d: + """Wrapper for JAX BatchNorm for 2D inputs.""" + + def __new__(cls, num_features, *args, **kwargs): + return nnx.BatchNorm(num_features, *args, **kwargs) + + +class LayerNorm: + """Wrapper for JAX LayerNorm.""" + + def __new__(cls, *args, **kwargs): + return nnx.LayerNorm(*args, **kwargs) diff --git a/illia/nn/jax/pooling.py b/illia/nn/jax/pooling.py new file mode 100644 index 0000000..78e0973 --- /dev/null +++ b/illia/nn/jax/pooling.py @@ -0,0 +1,35 @@ +"""JAX pooling layer wrappers.""" + +# Standard libraries +from functools import partial + +# 3pps +import flax.nnx as nnx + + +class MaxPool1d: + """Wrapper for JAX max_pool with 1D window.""" + + def __new__(cls, *args, **kwargs): + return partial(nnx.max_pool, *args, **kwargs) + + +class MaxPool2d: + """Wrapper for JAX max_pool with 2D window.""" + + def __new__(cls, *args, **kwargs): + return partial(nnx.max_pool, *args, **kwargs) + + +class AvgPool1d: + """Wrapper for JAX avg_pool with 1D window.""" + + def __new__(cls, *args, **kwargs): + return partial(nnx.avg_pool, *args, **kwargs) + + +class AvgPool2d: + """Wrapper for JAX avg_pool with 2D window.""" + + def __new__(cls, *args, **kwargs): + return partial(nnx.avg_pool, *args, **kwargs) diff --git a/illia/nn/jax/regularization.py b/illia/nn/jax/regularization.py new file mode 100644 index 0000000..54da00e --- /dev/null +++ b/illia/nn/jax/regularization.py @@ -0,0 +1,11 @@ +"""JAX regularization layer wrappers.""" + +# 3pps +from flax import nnx + + +class Dropout: + """Wrapper for JAX Dropout.""" + + def __new__(cls, *args, **kwargs): + return nnx.Dropout(*args, **kwargs) diff --git a/illia/nn/tf/__init__.py b/illia/nn/tf/__init__.py index 11edd1a..f69c20e 100644 --- a/illia/nn/tf/__init__.py +++ b/illia/nn/tf/__init__.py @@ -5,19 +5,46 @@ """ # Own modules +from illia.nn.tf.activation import GELU, LeakyReLU, ReLU, Sigmoid, Tanh from illia.nn.tf.base import BayesianModule from illia.nn.tf.conv1d import Conv1d from illia.nn.tf.conv2d import Conv2d from illia.nn.tf.embedding import Embedding from illia.nn.tf.linear import Linear from illia.nn.tf.lstm import LSTM +from illia.nn.tf.normalization import BatchNorm1d, BatchNorm2d, LayerNorm +from illia.nn.tf.pooling import ( + AdaptiveAvgPool2d, + AvgPool1d, + AvgPool2d, + MaxPool1d, + MaxPool2d, +) +from illia.nn.tf.regularization import Dropout, Dropout2d +from illia.nn.tf.utility import Flatten __all__: list[str] = [ + "AdaptiveAvgPool2d", + "AvgPool1d", + "AvgPool2d", + "BatchNorm1d", + "BatchNorm2d", "BayesianModule", "Conv1d", "Conv2d", + "Dropout", + "Dropout2d", "Embedding", - "Linear", + "Flatten", + "GELU", "LSTM", + "LayerNorm", + "LeakyReLU", + "Linear", + "MaxPool1d", + "MaxPool2d", + "ReLU", + "Sigmoid", + "Tanh", ] diff --git a/illia/nn/tf/activation.py b/illia/nn/tf/activation.py new file mode 100644 index 0000000..d4b3e50 --- /dev/null +++ b/illia/nn/tf/activation.py @@ -0,0 +1,39 @@ +"""TensorFlow activation layer wrappers.""" + +# 3pps +from tensorflow.keras import layers + + +class ReLU(layers.Activation): + """Wrapper for TensorFlow ReLU activation.""" + + def __init__(self, *args, **kwargs): + super().__init__("relu", *args, **kwargs) + + +class Sigmoid(layers.Activation): + """Wrapper for TensorFlow Sigmoid activation.""" + + def __init__(self, *args, **kwargs): + super().__init__("sigmoid", *args, **kwargs) + + +class Tanh(layers.Activation): + """Wrapper for TensorFlow Tanh activation.""" + + def __init__(self, *args, **kwargs): + super().__init__("tanh", *args, **kwargs) + + +class LeakyReLU(layers.Activation): + """Wrapper for TensorFlow LeakyReLU.""" + + def __init__(self, *args, **kwargs): + super().__init__("leaky_relu", *args, **kwargs) + + +class GELU(layers.Activation): + """Wrapper for TensorFlow GELU activation.""" + + def __init__(self, *args, **kwargs): + super().__init__("gelu", *args, **kwargs) diff --git a/illia/nn/tf/normalization.py b/illia/nn/tf/normalization.py new file mode 100644 index 0000000..970b2d7 --- /dev/null +++ b/illia/nn/tf/normalization.py @@ -0,0 +1,16 @@ +"""TensorFlow normalization layer wrappers.""" + +# 3pps +from tensorflow.keras import layers + + +class BatchNorm1d(layers.BatchNormalization): + """Wrapper for TensorFlow BatchNormalization (1D).""" + + +class BatchNorm2d(layers.BatchNormalization): + """Wrapper for TensorFlow BatchNormalization (2D).""" + + +class LayerNorm(layers.LayerNormalization): + """Wrapper for TensorFlow LayerNormalization.""" diff --git a/illia/nn/tf/pooling.py b/illia/nn/tf/pooling.py new file mode 100644 index 0000000..ccd0ca5 --- /dev/null +++ b/illia/nn/tf/pooling.py @@ -0,0 +1,24 @@ +"""TensorFlow pooling layer wrappers.""" + +# 3pps +from tensorflow.keras import layers + + +class MaxPool1d(layers.MaxPooling1D): + """Wrapper for TensorFlow MaxPooling1D.""" + + +class MaxPool2d(layers.MaxPooling2D): + """Wrapper for TensorFlow MaxPooling2D.""" + + +class AvgPool1d(layers.AveragePooling1D): + """Wrapper for TensorFlow AveragePooling1D.""" + + +class AvgPool2d(layers.AveragePooling2D): + """Wrapper for TensorFlow AveragePooling2D.""" + + +class AdaptiveAvgPool2d(layers.GlobalAveragePooling2D): + """Wrapper for TensorFlow GlobalAveragePooling2D.""" diff --git a/illia/nn/tf/regularization.py b/illia/nn/tf/regularization.py new file mode 100644 index 0000000..61fee1f --- /dev/null +++ b/illia/nn/tf/regularization.py @@ -0,0 +1,12 @@ +"""TensorFlow regularization layer wrappers.""" + +# 3pps +from tensorflow.keras import layers + + +class Dropout(layers.Dropout): + """Wrapper for TensorFlow Dropout.""" + + +class Dropout2d(layers.SpatialDropout2D): + """Wrapper for TensorFlow SpatialDropout2D.""" diff --git a/illia/nn/tf/utility.py b/illia/nn/tf/utility.py new file mode 100644 index 0000000..d9f2499 --- /dev/null +++ b/illia/nn/tf/utility.py @@ -0,0 +1,8 @@ +"""TensorFlow utility layer wrappers.""" + +# 3pps +from tensorflow.keras import layers + + +class Flatten(layers.Flatten): + """Wrapper for TensorFlow Flatten.""" diff --git a/illia/nn/torch/__init__.py b/illia/nn/torch/__init__.py index 3a2760c..a3d143f 100644 --- a/illia/nn/torch/__init__.py +++ b/illia/nn/torch/__init__.py @@ -5,19 +5,49 @@ """ # Own modules +from illia.nn.torch.activation import GELU, LeakyReLU, ReLU, Sigmoid, Tanh from illia.nn.torch.base import BayesianModule from illia.nn.torch.conv1d import Conv1d from illia.nn.torch.conv2d import Conv2d from illia.nn.torch.embedding import Embedding from illia.nn.torch.linear import Linear from illia.nn.torch.lstm import LSTM +from illia.nn.torch.normalization import BatchNorm1d, BatchNorm2d, LayerNorm +from illia.nn.torch.pooling import ( + AdaptiveAvgPool2d, + AdaptiveMaxPool2d, + AvgPool1d, + AvgPool2d, + MaxPool1d, + MaxPool2d, +) +from illia.nn.torch.regularization import Dropout, Dropout2d +from illia.nn.torch.utility import Flatten, Identity __all__: list[str] = [ + "AdaptiveAvgPool2d", + "AdaptiveMaxPool2d", + "AvgPool1d", + "AvgPool2d", + "BatchNorm1d", + "BatchNorm2d", "BayesianModule", "Conv1d", "Conv2d", + "Dropout", + "Dropout2d", "Embedding", - "Linear", + "Flatten", + "GELU", + "Identity", "LSTM", + "LayerNorm", + "LeakyReLU", + "Linear", + "MaxPool1d", + "MaxPool2d", + "ReLU", + "Sigmoid", + "Tanh", ] diff --git a/illia/nn/torch/activation.py b/illia/nn/torch/activation.py new file mode 100644 index 0000000..8590551 --- /dev/null +++ b/illia/nn/torch/activation.py @@ -0,0 +1,24 @@ +"""PyTorch activation layer wrappers.""" + +# 3pps +from torch import nn + + +class ReLU(nn.ReLU): + """Wrapper for PyTorch ReLU.""" + + +class Sigmoid(nn.Sigmoid): + """Wrapper for PyTorch Sigmoid.""" + + +class Tanh(nn.Tanh): + """Wrapper for PyTorch Tanh.""" + + +class LeakyReLU(nn.LeakyReLU): + """Wrapper for PyTorch LeakyReLU.""" + + +class GELU(nn.GELU): + """Wrapper for PyTorch GELU.""" diff --git a/illia/nn/torch/normalization.py b/illia/nn/torch/normalization.py new file mode 100644 index 0000000..5b92256 --- /dev/null +++ b/illia/nn/torch/normalization.py @@ -0,0 +1,16 @@ +"""PyTorch normalization layer wrappers.""" + +# 3pps +import torch.nn as nn + + +class BatchNorm1d(nn.BatchNorm1d): + """Wrapper for PyTorch BatchNorm1d.""" + + +class BatchNorm2d(nn.BatchNorm2d): + """Wrapper for PyTorch BatchNorm2d.""" + + +class LayerNorm(nn.LayerNorm): + """Wrapper for PyTorch LayerNorm.""" diff --git a/illia/nn/torch/pooling.py b/illia/nn/torch/pooling.py new file mode 100644 index 0000000..76a9dbd --- /dev/null +++ b/illia/nn/torch/pooling.py @@ -0,0 +1,28 @@ +"""PyTorch pooling layer wrappers.""" + +# 3pps +import torch.nn as nn + + +class MaxPool1d(nn.MaxPool1d): + """Wrapper for PyTorch MaxPool1d.""" + + +class MaxPool2d(nn.MaxPool2d): + """Wrapper for PyTorch MaxPool2d.""" + + +class AvgPool1d(nn.AvgPool1d): + """Wrapper for PyTorch AvgPool1d.""" + + +class AvgPool2d(nn.AvgPool2d): + """Wrapper for PyTorch AvgPool2d.""" + + +class AdaptiveAvgPool2d(nn.AdaptiveAvgPool2d): + """Wrapper for PyTorch AdaptiveAvgPool2d.""" + + +class AdaptiveMaxPool2d(nn.AdaptiveMaxPool2d): + """Wrapper for PyTorch AdaptiveMaxPool2d.""" diff --git a/illia/nn/torch/regularization.py b/illia/nn/torch/regularization.py new file mode 100644 index 0000000..a3b335c --- /dev/null +++ b/illia/nn/torch/regularization.py @@ -0,0 +1,12 @@ +"""PyTorch regularization layer wrappers.""" + +# 3pps +import torch.nn as nn + + +class Dropout(nn.Dropout): + """Wrapper for PyTorch Dropout.""" + + +class Dropout2d(nn.Dropout2d): + """Wrapper for PyTorch Dropout2d.""" diff --git a/illia/nn/torch/utility.py b/illia/nn/torch/utility.py new file mode 100644 index 0000000..9879742 --- /dev/null +++ b/illia/nn/torch/utility.py @@ -0,0 +1,12 @@ +"""PyTorch utility layer wrappers.""" + +# 3pps +import torch.nn as nn + + +class Flatten(nn.Flatten): + """Wrapper for PyTorch Flatten.""" + + +class Identity(nn.Identity): + """Wrapper for PyTorch Identity.""" diff --git a/illia/support.py b/illia/support.py index 08448b7..1127354 100644 --- a/illia/support.py +++ b/illia/support.py @@ -48,28 +48,103 @@ # Dictionary describing the layers and capabilities supported by each backend BACKEND_CAPABILITIES: Final[dict[str, dict[str, set[str]]]] = { "torch": { - "nn": {"BayesianModule", "Conv1d", "Conv2d", "Embedding", "Linear", "LSTM"}, - "distributions": {"DistributionModule", "GaussianDistribution"}, - "losses": { - "KLDivergenceLoss", - "ELBOLoss", + "nn": { + # Bayesian layers + "BayesianModule", + "Conv1d", + "Conv2d", + "Embedding", + "Linear", + "LSTM", + # Pooling layers + "MaxPool1d", + "MaxPool2d", + "AvgPool1d", + "AvgPool2d", + "AdaptiveAvgPool2d", + "AdaptiveMaxPool2d", + # Activation layers + "ReLU", + "Sigmoid", + "Tanh", + "LeakyReLU", + "GELU", + # Normalization layers + "BatchNorm1d", + "BatchNorm2d", + "LayerNorm", + # Regularization layers + "Dropout", + "Dropout2d", + # Utility layers + "Flatten", + "Identity", }, + "distributions": {"DistributionModule", "GaussianDistribution"}, + "losses": {"KLDivergenceLoss", "ELBOLoss"}, }, "tf": { - "nn": {"BayesianModule", "Conv1d", "Conv2d", "Embedding", "Linear", "LSTM"}, - "distributions": {"DistributionModule", "GaussianDistribution"}, - "losses": { - "KLDivergenceLoss", - "ELBOLoss", + "nn": { + # Bayesian layers + "BayesianModule", + "Conv1d", + "Conv2d", + "Embedding", + "Linear", + "LSTM", + # Pooling layers + "MaxPool1d", + "MaxPool2d", + "AvgPool1d", + "AvgPool2d", + "AdaptiveAvgPool2d", + # Activation layers + "ReLU", + "Sigmoid", + "Tanh", + "LeakyReLU", + "GELU", + # Normalization layers + "BatchNorm1d", + "BatchNorm2d", + "LayerNorm", + # Regularization layers + "Dropout", + "Dropout2d", + # Utility layers + "Flatten", }, + "distributions": {"DistributionModule", "GaussianDistribution"}, + "losses": {"KLDivergenceLoss", "ELBOLoss"}, }, "jax": { - "nn": {"BayesianModule", "Conv1d", "Conv2d", "Embedding", "Linear", "LSTM"}, - "distributions": {"DistributionModule", "GaussianDistribution"}, - "losses": { - "KLDivergenceLoss", - "ELBOLoss", + "nn": { + # Bayesian layers + "BayesianModule", + "Conv1d", + "Conv2d", + "Embedding", + "Linear", + "LSTM", + # Pooling layers + "MaxPool1d", + "MaxPool2d", + "AvgPool1d", + "AvgPool2d", + # Activation layers + "ReLU", + "Sigmoid", + "Tanh", + "GELU", + # Normalization layers + "BatchNorm1d", + "BatchNorm2d", + "LayerNorm", + # Regularization layers + "Dropout", }, + "distributions": {"DistributionModule", "GaussianDistribution"}, + "losses": {"KLDivergenceLoss", "ELBOLoss"}, }, "pyg": { "nn": {"CGConv"}, diff --git a/tests/jax/nn/conftest.py b/tests/jax/nn/conftest.py index 52b67c3..072154a 100644 --- a/tests/jax/nn/conftest.py +++ b/tests/jax/nn/conftest.py @@ -21,6 +21,12 @@ from illia.nn import LSTM, Conv1d, Conv2d, Embedding, Linear +@pytest.fixture(scope="function") +def rngs_fixture(): + """Provide consistent RNG for JAX tests.""" + return nnx.Rngs(42) + + @pytest.fixture( params=[ (32, 30, 20, None, None), diff --git a/tests/jax/nn/test_non_parametric.py b/tests/jax/nn/test_non_parametric.py new file mode 100644 index 0000000..864b245 --- /dev/null +++ b/tests/jax/nn/test_non_parametric.py @@ -0,0 +1,108 @@ +""" +This module contains the tests for the Non-Bayesian layers Wrapper. +""" + +# Standard libraries +import os + + +# Change Illia Backend +os.environ["ILLIA_BACKEND"] = "jax" + +# 3pps +import jax.numpy as jnp +import pytest + + +class TestNonParametricLayers: + """ + This class tests the non-parametric layers integration. + """ + + @pytest.mark.order(1) + @pytest.mark.parametrize( + "layer_name,input_shape,window_shape", + [ + ("MaxPool1d", (32, 16, 28), (2,)), + ("MaxPool2d", (32, 16, 28, 28), (2, 2)), + ("AvgPool1d", (32, 16, 28), (2,)), + ("AvgPool2d", (32, 16, 28, 28), (2, 2)), + ], + ) + def test_pooling( + self, layer_name: str, input_shape: tuple, window_shape: tuple + ) -> None: + """ + Test pooling functions. + + Args: + layer_name: Name of the pooling function to test. + input_shape: Shape of the input tensor for testing. + window_shape: Shape of the pooling window. + """ + pool_fn = getattr(__import__("illia.nn", fromlist=[layer_name]), layer_name)() + inputs = jnp.ones(input_shape) + output = pool_fn(inputs, window_shape=window_shape) + assert isinstance(output, jnp.ndarray) + assert output.dtype == inputs.dtype + + @pytest.mark.order(2) + @pytest.mark.parametrize("layer_name", ["ReLU", "Sigmoid", "Tanh", "GELU"]) + def test_activation(self, layer_name: str) -> None: + """ + Test activation functions. + + Args: + layer_name: Name of the activation function to test. + """ + activation_fn = getattr( + __import__("illia.nn", fromlist=[layer_name]), layer_name + )() + inputs = jnp.ones((32, 16, 28, 28)) + output = activation_fn(inputs) + assert isinstance(output, jnp.ndarray) + assert tuple(output.shape) == tuple(inputs.shape) + assert output.dtype == inputs.dtype + + @pytest.mark.order(3) + def test_regularization(self, rngs_fixture) -> None: + """ + Test regularization layers. + + Args: + rngs_fixture: JAX RNG fixture for consistent random number generation. + """ + Dropout = getattr(__import__("illia.nn", fromlist=["Dropout"]), "Dropout") + layer_instance = Dropout(rate=0.5, rngs=rngs_fixture) + assert layer_instance is not None + assert callable(layer_instance) + + @pytest.mark.order(4) + def test_normalization(self, rngs_fixture) -> None: + """ + Test normalization layers. + + Args: + rngs_fixture: JAX RNG fixture for consistent random number generation. + """ + # Test BatchNorm1d + BatchNorm1d = getattr( + __import__("illia.nn", fromlist=["BatchNorm1d"]), "BatchNorm1d" + ) + layer1d = BatchNorm1d(num_features=16, rngs=rngs_fixture) + assert layer1d is not None + assert callable(layer1d) + + # Test BatchNorm2d + BatchNorm2d = getattr( + __import__("illia.nn", fromlist=["BatchNorm2d"]), "BatchNorm2d" + ) + layer2d = BatchNorm2d(num_features=16, rngs=rngs_fixture) + assert layer2d is not None + assert callable(layer2d) + + # Test LayerNorm + LayerNorm = getattr(__import__("illia.nn", fromlist=["LayerNorm"]), "LayerNorm") + layer_norm = LayerNorm(num_features=16, rngs=rngs_fixture) + assert layer_norm is not None + assert callable(layer_norm) diff --git a/tests/jax/nn/test_serialize.py b/tests/jax/nn/test_serialize.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/tf/nn/conftest.py b/tests/tf/nn/conftest.py index df26f46..e2124b2 100644 --- a/tests/tf/nn/conftest.py +++ b/tests/tf/nn/conftest.py @@ -20,6 +20,13 @@ from illia.nn import LSTM, Conv1d, Conv2d, Embedding, Linear +@pytest.fixture(scope="session", autouse=True) +def set_random_seeds(): + """Set random seeds for reproducibility.""" + tf.keras.utils.set_random_seed(42) + tf.config.experimental.enable_op_determinism() + + @pytest.fixture( params=[ (32, 30, 20, None, None), diff --git a/tests/tf/nn/test_non_parametric.py b/tests/tf/nn/test_non_parametric.py new file mode 100644 index 0000000..c38e502 --- /dev/null +++ b/tests/tf/nn/test_non_parametric.py @@ -0,0 +1,132 @@ +""" +This module contains the tests for the Non-Bayesian layers Wrapper. +""" + +# Standard libraries +import os + + +# Change Illia Backend +os.environ["ILLIA_BACKEND"] = "tf" + +# 3pps +import pytest +import tensorflow as tf + + +class TestNonParametricLayers: + """ + This class tests the non-parametric layers integration. + """ + + @pytest.mark.order(1) + @pytest.mark.parametrize( + "layer_name,kwargs,input_shape,expected_shape", + [ + ("MaxPool2d", {"pool_size": 2}, (32, 28, 28, 16), (32, 14, 14, 16)), + ("AvgPool2d", {"pool_size": 2}, (32, 28, 28, 16), (32, 14, 14, 16)), + ("MaxPool1d", {"pool_size": 2}, (32, 64, 16), (32, 32, 16)), + ("AvgPool1d", {"pool_size": 2}, (32, 64, 16), (32, 32, 16)), + ("AdaptiveAvgPool2d", {}, (32, 28, 28, 16), (32, 16)), + ], + ) + def test_pooling( + self, layer_name: str, kwargs: dict, input_shape: tuple, expected_shape: tuple + ) -> None: + """ + Test pooling layers. + + Args: + layer_name: Name of the pooling layer to test. + kwargs: Keyword arguments for layer initialization. + input_shape: Shape of the input tensor for testing. + expected_shape: Expected shape of the output tensor. + """ + layer = getattr(__import__("illia.nn", fromlist=[layer_name]), layer_name)( + **kwargs + ) + inputs = tf.random.uniform(input_shape) + output = layer(inputs) + assert isinstance(output, tf.Tensor) + assert tuple(output.shape) == expected_shape + assert output.dtype == inputs.dtype + + @pytest.mark.order(2) + @pytest.mark.parametrize( + "layer_name", ["ReLU", "Sigmoid", "Tanh", "LeakyReLU", "GELU"] + ) + def test_activation(self, layer_name: str) -> None: + """ + Test activation layers. + + Args: + layer_name: Name of the activation layer to test. + """ + layer = getattr(__import__("illia.nn", fromlist=[layer_name]), layer_name)() + inputs = tf.random.uniform((32, 28, 28, 16)) + output = layer(inputs) + assert isinstance(output, tf.Tensor) + assert tuple(output.shape) == tuple(inputs.shape) + assert output.dtype == inputs.dtype + + @pytest.mark.order(3) + @pytest.mark.parametrize( + "layer_name,input_shape", + [ + ("BatchNorm1d", (32, 64, 16)), + ("BatchNorm2d", (32, 28, 28, 16)), + ("LayerNorm", (32, 28, 28, 16)), + ], + ) + def test_normalization(self, layer_name: str, input_shape: tuple) -> None: + """ + Test normalization layers. + + Args: + layer_name: Name of the normalization layer to test. + input_shape: Shape of the input tensor for testing. + """ + layer = getattr(__import__("illia.nn", fromlist=[layer_name]), layer_name)() + inputs = tf.random.normal(input_shape, mean=5.0, stddev=2.0) + output = layer(inputs, training=True) + assert isinstance(output, tf.Tensor) + assert tuple(output.shape) == tuple(inputs.shape) + assert output.dtype == inputs.dtype + + @pytest.mark.order(4) + @pytest.mark.parametrize( + "layer_name,rate,input_shape", + [ + ("Dropout", 0.5, (32, 28, 28, 16)), + ("Dropout2d", 0.5, (32, 28, 28, 16)), + ], + ) + def test_regularization( + self, layer_name: str, rate: float, input_shape: tuple + ) -> None: + """ + Test regularization layers. + + Args: + layer_name: Name of the regularization layer to test. + rate: Dropout rate for the regularization layer. + input_shape: Shape of the input tensor for testing. + """ + layer = getattr(__import__("illia.nn", fromlist=[layer_name]), layer_name)(rate) + inputs = tf.random.uniform(input_shape) + output = layer(inputs, training=True) + assert isinstance(output, tf.Tensor) + assert tuple(output.shape) == tuple(inputs.shape) + assert output.dtype == inputs.dtype + + @pytest.mark.order(5) + def test_flatten(self) -> None: + """ + Test Flatten utility layer. + """ + layer = getattr(__import__("illia.nn", fromlist=["Flatten"]), "Flatten")() + inputs = tf.random.uniform((32, 28, 28, 16)) + output = layer(inputs) + assert isinstance(output, tf.Tensor) + assert tuple(output.shape) == (32, 28 * 28 * 16) + assert output.dtype == inputs.dtype diff --git a/tests/tf/nn/test_serialize.py b/tests/tf/nn/test_serialize.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/torch/nn/conftest.py b/tests/torch/nn/conftest.py index d6b6f37..81c176c 100644 --- a/tests/torch/nn/conftest.py +++ b/tests/torch/nn/conftest.py @@ -21,6 +21,15 @@ from tests.torch.nn.utils import BayesianComposedModel, ComposedModel +@pytest.fixture(scope="session", autouse=True) +def set_random_seeds(): + """Set random seeds for reproducibility.""" + torch.manual_seed(42) + torch.cuda.manual_seed_all(42) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + @pytest.fixture( params=[ (32, 30, 20, None, None), diff --git a/tests/torch/nn/test_non_parametric.py b/tests/torch/nn/test_non_parametric.py new file mode 100644 index 0000000..6115440 --- /dev/null +++ b/tests/torch/nn/test_non_parametric.py @@ -0,0 +1,164 @@ +""" +This module contains the tests for the Non-Bayesian layers Wrapper. +""" + +# Standard libraries +import os + + +# Change Illia Backend +os.environ["ILLIA_BACKEND"] = "torch" + +# 3pps +import pytest +import torch + + +class TestNonParametricLayers: + """ + This class tests the non-parametric layers integration. + """ + + @pytest.mark.order(1) + @pytest.mark.parametrize( + "layer_name,kwargs,input_shape,expected_shape", + [ + ("MaxPool2d", {"kernel_size": 2}, (32, 16, 28, 28), (32, 16, 14, 14)), + ("AvgPool2d", {"kernel_size": 2}, (32, 16, 28, 28), (32, 16, 14, 14)), + ("MaxPool1d", {"kernel_size": 2}, (32, 16, 64), (32, 16, 32)), + ("AvgPool1d", {"kernel_size": 2}, (32, 16, 64), (32, 16, 32)), + ( + "AdaptiveAvgPool2d", + {"output_size": (1, 1)}, + (32, 16, 28, 28), + (32, 16, 1, 1), + ), + ( + "AdaptiveMaxPool2d", + {"output_size": (1, 1)}, + (32, 16, 28, 28), + (32, 16, 1, 1), + ), + ], + ) + def test_pooling( + self, layer_name: str, kwargs: dict, input_shape: tuple, expected_shape: tuple + ) -> None: + """ + Test pooling layers. + + Args: + layer_name: Name of the pooling layer to test. + kwargs: Keyword arguments for layer initialization. + input_shape: Shape of the input tensor for testing. + expected_shape: Expected shape of the output tensor. + """ + layer = getattr(__import__("illia.nn", fromlist=[layer_name]), layer_name)( + **kwargs + ) + inputs = torch.rand(input_shape) + output = layer(inputs) + assert isinstance(output, torch.Tensor) + assert tuple(output.shape) == expected_shape + assert output.dtype == inputs.dtype + + @pytest.mark.order(2) + @pytest.mark.parametrize( + "layer_name", ["ReLU", "Sigmoid", "Tanh", "LeakyReLU", "GELU"] + ) + def test_activation(self, layer_name: str) -> None: + """ + Test activation layers. + + Args: + layer_name: Name of the activation layer to test. + """ + layer = getattr(__import__("illia.nn", fromlist=[layer_name]), layer_name)() + inputs = torch.rand((32, 16, 28, 28)) + output = layer(inputs) + assert isinstance(output, torch.Tensor) + assert tuple(output.shape) == tuple(inputs.shape) + assert output.dtype == inputs.dtype + + @pytest.mark.order(3) + @pytest.mark.parametrize( + "layer_name,kwargs,input_shape", + [ + ("BatchNorm1d", {"num_features": 16}, (32, 16, 64)), + ("BatchNorm2d", {"num_features": 16}, (32, 16, 28, 28)), + ("LayerNorm", {"normalized_shape": [16, 28, 28]}, (32, 16, 28, 28)), + ], + ) + def test_normalization( + self, layer_name: str, kwargs: dict, input_shape: tuple + ) -> None: + """ + Test normalization layers. + + Args: + layer_name: Name of the normalization layer to test. + kwargs: Keyword arguments for layer initialization. + input_shape: Shape of the input tensor for testing. + """ + layer = getattr(__import__("illia.nn", fromlist=[layer_name]), layer_name)( + **kwargs + ) + layer.train() + inputs = torch.randn(input_shape) * 2.0 + 5.0 + output = layer(inputs) + assert isinstance(output, torch.Tensor) + assert tuple(output.shape) == tuple(inputs.shape) + assert output.dtype == inputs.dtype + + @pytest.mark.order(4) + @pytest.mark.parametrize( + "layer_name,rate,input_shape", + [ + ("Dropout", 0.5, (32, 16, 28, 28)), + ("Dropout2d", 0.5, (32, 16, 28, 28)), + ], + ) + def test_regularization( + self, layer_name: str, rate: float, input_shape: tuple + ) -> None: + """ + Test regularization layers. + + Args: + layer_name: Name of the regularization layer to test. + rate: Dropout rate for the regularization layer. + input_shape: Shape of the input tensor for testing. + """ + layer = getattr(__import__("illia.nn", fromlist=[layer_name]), layer_name)(rate) + layer.train() + inputs = torch.rand(input_shape) + output = layer(inputs) + assert isinstance(output, torch.Tensor) + assert tuple(output.shape) == tuple(inputs.shape) + assert output.dtype == inputs.dtype + + @pytest.mark.order(5) + @pytest.mark.parametrize( + "layer_name,input_shape,expected_shape", + [ + ("Flatten", (32, 16, 28, 28), (32, 16 * 28 * 28)), + ("Identity", (32, 16, 28, 28), (32, 16, 28, 28)), + ], + ) + def test_utility( + self, layer_name: str, input_shape: tuple, expected_shape: tuple + ) -> None: + """ + Test utility layers. + + Args: + layer_name: Name of the utility layer to test. + input_shape: Shape of the input tensor for testing. + expected_shape: Expected shape of the output tensor. + """ + layer = getattr(__import__("illia.nn", fromlist=[layer_name]), layer_name)() + inputs = torch.rand(input_shape) + output = layer(inputs) + assert isinstance(output, torch.Tensor) + assert tuple(output.shape) == expected_shape + assert output.dtype == inputs.dtype diff --git a/tests/torch/nn/test_serialize.py b/tests/torch/nn/test_serialize.py new file mode 100644 index 0000000..e69de29