stax layer with costume activation function

[ThatPhysicsGuy]

I'm trying to implement different activation functions for a ML project I'm working on. I've tried to use leaky_relu but it gives me a << TypeError: 'CompiledFunction' object is not iterable >> . So I thought to write my own leaky_relu and add it as customized activation for the layer. I saw that with Tensorflow is doable but I can't find any documentation about a layer with customizable activation function for STAX.

Here's the model I'm implementing:

stax.serial(
stax.Dense(28,),stax.leaky_relu,
stax.Dense(128),stax.leaky_relu,
stax.Dense(1)}