original_docstring: 
      Execute layers in parallel then reduce them back together
      
docstring: Execute layers in parallel then reduce them back together
docstring_tokens: Execute layers in parallel then reduce them back together
code: def parallel(*layers, reduce_fun=reduce_add):
      """
      Execute layers in parallel then reduce them back together
      """
      nlayers = len(layers)
      init_funs, apply_funs = zip(*layers)
  
      def init_fun(rng, input_shape):
          rngs = random.split(rng, num=nlayers)
          params, out_shape = [], None
          for r, layer_init in zip(rngs, init_funs):
              o, p = layer_init(r, input_shape)
              params.append(p)
              out_shape = o if out_shape is None else out_shape
              if o != out_shape:
                  raise ValueError("Layers don't have same output shapes; can't reduce")
          return out_shape, params
  
      def apply_fun(params, inputs, **kwargs):
          rng = kwargs.pop("rng", None)
          rngs = random.split(rng, nlayers) if rng is not None else (None,) * nlayers
          outputs = [
              f(p, inputs, rng=r, **kwargs) for f, p, r in zip(apply_funs, params, rngs)
          ]
          return reduce_fun(*outputs)
  
      return init_fun, apply_fun
code_tokens: def parallel ( * layers , reduce_fun = reduce_add ) : nlayers = len ( layers ) init_funs , apply_funs = zip ( * layers ) def init_fun ( rng , input_shape ) : rngs = random . split ( rng , num = nlayers ) params , out_shape = [ ] , None for r , layer_init in zip ( rngs , init_funs ) : o , p = layer_init ( r , input_shape ) params . append ( p ) out_shape = o if out_shape is None else out_shape if o != out_shape : raise ValueError ( "Layers don't have same output shapes; can't reduce" ) return out_shape , params def apply_fun ( params , inputs , ** kwargs ) : rng = kwargs . pop ( "rng" , None ) rngs = random . split ( rng , nlayers ) if rng is not None else ( None , ) * nlayers outputs = [ f ( p , inputs , rng = r , ** kwargs ) for f , p , r in zip ( apply_funs , params , rngs ) ] return reduce_fun ( * outputs ) return init_fun , apply_fun
short_docstring: Execute layers in parallel then reduce them back together
short_docstring_tokens: Execute layers in parallel then reduce them back together
comment: """
      Execute layers in parallel then reduce them back together
      """