Apply jnp.exp to output of a model

[averageFlaxUser]

I have a NN that outputs mu and log_sigma, essentially parameters of a normal distribution. To avoid negative values, I want to apply sigma = jnp.exp(log_sigma) before doing:

mu + jax.random.normal(gaussian_key, shape=(1, 100)) * sigma

But I get a warning from Jax at jnp.exp(log_sigma):

Argument of type "Any | FrozenVariableDict" cannot be assigned to parameter of type "ArrayLike"

How to deal with this?

My model is defined using flax as such:

class Model(nn.Module):
    training: bool    
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(features=10)(x)
        x = nn.relu(x)
        x = nn.Dense(features=1024)(x)
        x = nn.relu(x)
        x = nn.Dropout(rate=0.50, deterministic=not self.training)(x)
        x = nn.Dense(features=2)(x)
        return x

[PhilipVinc]

Can't you simply redefine your model to be

class Model(nn.Module):
    training: bool    
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(features=10)(x)
        ...
        x = nn.Dense(features=2)(x)
        mu, sigma = x[...,0], x[...,1]
        return mu, jnp.exp(sigma)