test_activations.py

"""
Tests for Activation Functions
"""

import pytest
import numpy as np

try:
    import _kortexdl_core as bd
    BINDINGS_AVAILABLE = True
except ImportError:
    BINDINGS_AVAILABLE = False
    bd = None

pytestmark = pytest.mark.requires_build

@pytest.mark.parametrize("activation_type", [
    bd.ActivationType.ReLU,
    bd.ActivationType.LeakyReLU,
    bd.ActivationType.Sigmoid,
    bd.ActivationType.Tanh,
    bd.ActivationType.Linear,
    bd.ActivationType.ELU,
    bd.ActivationType.Swish,
    # bd.ActivationType.Softmax  # Softmax usually requires special handling (output sum=1), might fail basic range check loop if applied elementwise? 
    # Activation.cpp usually applies Softmax row-wise. 
    # Let's keep Softmax but adapt test if needed.
    bd.ActivationType.Softmax
])
def test_activation_forward(activation_type):
    """Test verification of activation forward pass (basic properties)"""
    # Create layer 1->1
    layer = bd.Layer(1, 1, activation_type)
    # W=1, b=0
    layer.set_parameters([[1.0]], [0.0])
    
    inputs = np.array([-1.0, 0.0, 1.0], dtype=np.float32)
    # Using 3 batch size
    output = layer.forward(inputs, 3)
    output = np.array(output)
    
    if activation_type == bd.ActivationType.ReLU:
        assert np.all(output >= 0)
        assert output[0] == 0 # ReLU(-1) = 0
    elif activation_type == bd.ActivationType.Sigmoid:
        assert np.all((output >= 0) & (output <= 1))
    elif activation_type == bd.ActivationType.Tanh:
        assert np.all((output >= -1) & (output <= 1))
    elif activation_type == bd.ActivationType.Linear:
        np.testing.assert_almost_equal(output, inputs)
        
    # Check for NaN
    assert not np.any(np.isnan(output))

@pytest.mark.parametrize("activation_type", [
    bd.ActivationType.ReLU,
    bd.ActivationType.Sigmoid,
    bd.ActivationType.Tanh
])
def test_activation_gradient(activation_type):
    """Test activation gradients are computed (non-zero)"""
    layer = bd.Layer(1, 1, activation_type)
    layer.set_parameters([[1.0]], [0.0])
    inputs = np.array([0.5, 0.1], dtype=np.float32) # Avoid 0 (ReLU) and symmetric input (Cancellation)
    
    # Must forward first!
    layer.forward(inputs, 2)
    
    grad_out = np.ones_like(inputs)
    layer.backward(inputs, grad_out, 0.0, 2)
    
    grad_w = layer.get_grad_weights()
    # Gradient shouldn't be zero/nan
    assert not np.any(np.isnan(grad_w))
    # For these inputs, grad should be non-zero
    assert np.sum(np.abs(grad_w)) > 1e-10