code: def Flatten():
      """Layer construction function for flattening all but the leading dim."""
  
      def init_fun(rng, input_shape):
          output_shape = input_shape[0], functools.reduce(
              op.mul, input_shape[1:], 1
          )
          return output_shape, ()
  
      def apply_fun(params, inputs, **kwargs):
          return jnp.reshape(inputs, (inputs.shape[0], -1))
  
      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 8 for calling a global function with example usage: Flatten() and returns: init_fun,  apply_fun, output_shape,  , jnp