How can we define different activation functions for different components?

[blackblitz]

I would like to define different activation functions for different components. For example, I have two nodes in a layer and would like to have identity activation for the first node and square activation for the second node. The following is my attempt.

import flax.linen as nn
import jax.numpy as jnp

class MyFunc(nn.Module):

    @nn.compact
    def __call__(self, x):
        x = nn.Dense(2)(x)
        x = jnp.array([x[0], x[1] ** 2])
        return x

This works fine if I call model.apply on a one-dimensional x, but if I perform batch apply with two-dimensional x, it does not give the correct shape. For example, if I input x with shape (100, 1), it should output y with shape (100, 2), but the output actually has shape (2, 2). How should I properly define the activation functions?

[cgarciae]

Use jnp.stack.