code: def FanInConcat(axis=-1):
      """Layer construction function for a fan-in concatenation layer."""
  
      def init_fun(rng, input_shape):
          ax = axis % len(input_shape[0])
          concat_size = sum(shape[ax] for shape in input_shape)
          out_shape = (
              input_shape[0][:ax]
              + (concat_size,)
              + input_shape[0][ax + 1 :]
          )
          return out_shape, ()
  
      def apply_fun(params, inputs, **kwargs):
          return jnp.concatenate(inputs, axis)
  
      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 12 for calling a global function with example usage: FanInConcat(axis) and returns: init_fun,  apply_fun, out_shape,  , jnp