code: def elementwise(fun, **fun_kwargs):
      """Layer that applies a scalar function elementwise on its inputs."""
      init_fun = lambda rng, input_shape: (input_shape, ())
      apply_fun = lambda params, inputs, **kwargs: fun(
          inputs, **fun_kwargs
      )
      return init_fun, apply_fun
  
imports: import functools
  import operator
  from jax import lax
  from jax import random
  import jax.numpy
  from jax.nn import relu, log_softmax, softmax, softplus, sigmoid, elu, leaky_relu, selu, gelu, standardize
  from jax.nn.initializers import glorot_normal, normal, ones, zeros
label_desc: Use image node_id 5 for calling a global function with example usage: elementwise(fun) and returns: init_fun,  apply_fun