PReLU Activation

[enver1323]

Hi @patrick-kidger, I am a big fan of Equinox, thank you for publishing this framework. I am currently working on a project and decided to give a PReLU a try, Therefore, I was wondering, if it was possible to add PReLU Activation function to a some module like nn.activations.

My current implementation, is as follows:

class PReLU(eqx.Module):
    negative_slope: jaxtyping.Array

    def __init__(self, alpha):
        self.negative_slope = jnp.array((alpha,))

    def __call__(self, x):
        return jax.numpy.where(x >= 0, x, self.negative_slope * x)

[patrick-kidger]

Your implementation looks good to me. I'd be happy to accept a PR adding this to Equinox.

[enver1323]

Thank you, @patrick-kidger! I added the PR with the implementation and docs in the PR #249 .

[patrick-kidger]

Closing since this now exists on the dev branch, which will shortly be merged to main.