There are many cases where I want to setup a flexible neural network parameterization where I try out different "activation" functions and final layers from the CLI. A convenient design for this is to use jsonargparse or alternatives and do things like

# Workaround: https://github.com/omni-us/jsonargparse/issues/619#issuecomment-2466451720
nnx.softplus.__deepcopy__ = MethodType(lambda self, memo: self, nnx.softplus)

def my_model(
    width: int = 64,
    depth: int = 4,
    activation: Callable = nnx.relu,
    final_activation: Callable = nnx.softplus,
    ...): 
    # use the settings to create a neural network
if __name__ == "__main__":
    jsonargparse.CLI(my_model)

Then I can do things like python my_model.py --depth=3 --activation=nnx.tanh

However, a common thing I need to do is swap out different settings for an identity function. But this isn't available in jax or nnx. What I would love to do is try things like python my_model.py --depth=3 --final_activation=nnx.identity

So, could we add this to the nnx activation functions? Something trivial like

@jax.jit
def identity(x: ArrayLike) -> Array:
  r"""Identity activation function.

  Returns the argument unmodified.

  Args:
    x : input array

  Returns:
    The argument `x` unmodified.

  Examples:
    >>> nnx.identity(jax.numpy.array([-2., -1., -0.5, 0, 0.5, 1., 2.]))
    Array([-2. , -1. , -0.5, 0. , 0.5, 1. , 2. ], dtype=float32)

  """  
  return x

If so, I can create a PR if you tell me which files to add it to. It seemed like https://github.com/google/flax/blob/main/flax/nnx/nn/activations.py is a natural place. If a test is required, tell me which file it makes the most sense in?