original_docstring: 
      Creates init and apply functions for a fully-connected NN with a Gaussian
      likelihood.
  
      :param output_dim: Number of output dimensions
      :param first_layer_factor: Scaling factor for first dense layer initialization.
      
docstring: Creates init and apply functions for a fully-connected NN with a Gaussian
  likelihood.
docstring_tokens: Creates init and apply functions for a fully - connected NN with a Gaussian likelihood .
code: def fcnn(
      output_dim: int,
      layers=2,
      units=256,
      skips=False,
      output_activation=None,
      first_layer_factor=None,
      parallel_with=None,
  ):
      """
      Creates init and apply functions for a fully-connected NN with a Gaussian
      likelihood.
  
      :param output_dim: Number of output dimensions
      :param first_layer_factor: Scaling factor for first dense layer initialization.
      """
  
      first_layer_factor = 10.0 if first_layer_factor is None else first_layer_factor
  
      activation = stax.elementwise(np.sin)
      block = (
          skip_connect(stax.serial(Dense(units, W_init=siren_w_init()()), activation))
          if skips
          else stax.serial(Dense(units, W_init=siren_w_init()()), activation)
      )
  
      layer_list = [
          stax.serial(
              Dense(
                  units,
                  W_init=siren_w_init(factor=first_layer_factor)(),
                  b_init=jax.nn.initializers.normal(stddev=2.0 * math.pi),
              ),
              activation,
          ),
          *([block] * (layers - 1)),
          Dense(output_dim),  # No skips on the last one!
      ]
      if output_activation is not None:
          layer_list.append(output_activation)
      if parallel_with is None:
          _init_fun, _apply_fun = stax.serial(*layer_list)
      else:
          _init_fun, _apply_fun = parallel(stax.serial(*layer_list), parallel_with)
  
      def init_fun(rng, input_shape):
          output_shape, net_params = _init_fun(rng, input_shape)
          params = {"net": net_params, "raw_noise": np.array(-2.0)}
          return output_shape, params
  
      # Conform to API:
      @t_wrapper
      def apply_fun(params, rng, inputs):
          return _apply_fun(params["net"], inputs)
  
      @t_wrapper
      def gaussian_fun(params, rng, inputs):
          pred_mean = apply_fun(params, None, inputs)
          pred_std = params["noise"] * np.ones(pred_mean.shape)
          return pred_mean, pred_std
  
      @t_wrapper
      def loss_fun(
          params, rng, data, batch_size=None, n=None, loss_type="nlp", reduce="sum"
      ):
          """
          :param batch_size: How large a batch to subselect from the provided data
          :param n: The total size of the dataset (to multiply batch estimate by)
          """
          assert loss_type in ("nlp", "mse")
          inputs, targets = data
          n = inputs.shape[0] if n is None else n
          if batch_size is not None:
              rng, rng_batch = random.split(rng)
              i = random.permutation(rng_batch, n)[:batch_size]
              inputs, targets = inputs[i], targets[i]
          preds = apply_fun(params, rng, inputs).squeeze()
          mean_loss = (
              -norm.logpdf(targets.squeeze(), preds, params["noise"]).mean()
              if loss_type == "nlp"
              else np.power(targets.squeeze() - preds, 2).mean()
          )
          if reduce == "sum":
              loss = n * mean_loss
          elif reduce == "mean":
              loss = mean_loss
          return loss
  
      def sample_fun_fun(rng, params):
          def f(x):
              return apply_fun(params, rng, x)
  
          return f
  
      return {
          "init": init_fun,
          "apply": apply_fun,
          "gaussian": gaussian_fun,
          "loss": loss_fun,
          "sample_function": sample_fun_fun,
      }
code_tokens: def fcnn ( output_dim : int , layers = 2 , units = 256 , skips = False , output_activation = None , first_layer_factor = None , parallel_with = None , ) : first_layer_factor = 10.0 if first_layer_factor is None else first_layer_factor activation = stax . elementwise ( np . sin ) block = ( skip_connect ( stax . serial ( Dense ( units , W_init = siren_w_init ( ) ( ) ) , activation ) ) if skips else stax . serial ( Dense ( units , W_init = siren_w_init ( ) ( ) ) , activation ) ) layer_list = [ stax . serial ( Dense ( units , W_init = siren_w_init ( factor = first_layer_factor ) ( ) , b_init = jax . nn . initializers . normal ( stddev = 2.0 * math . pi ) , ) , activation , ) , * ( [ block ] * ( layers - 1 ) ) , Dense ( output_dim ) , ] if output_activation is not None : layer_list . append ( output_activation ) if parallel_with is None : _init_fun , _apply_fun = stax . serial ( * layer_list ) else : _init_fun , _apply_fun = parallel ( stax . serial ( * layer_list ) , parallel_with ) def init_fun ( rng , input_shape ) : output_shape , net_params = _init_fun ( rng , input_shape ) params = { "net" : net_params , "raw_noise" : np . array ( - 2.0 ) } return output_shape , params @ t_wrapper def apply_fun ( params , rng , inputs ) : return _apply_fun ( params [ "net" ] , inputs ) @ t_wrapper def gaussian_fun ( params , rng , inputs ) : pred_mean = apply_fun ( params , None , inputs ) pred_std = params [ "noise" ] * np . ones ( pred_mean . shape ) return pred_mean , pred_std @ t_wrapper def loss_fun ( params , rng , data , batch_size = None , n = None , loss_type = "nlp" , reduce = "sum" ) : """
          :param batch_size: How large a batch to subselect from the provided data
          :param n: The total size of the dataset (to multiply batch estimate by)
          """ assert loss_type in ( "nlp" , "mse" ) inputs , targets = data n = inputs . shape [ 0 ] if n is None else n if batch_size is not None : rng , rng_batch = random . split ( rng ) i = random . permutation ( rng_batch , n ) [ : batch_size ] inputs , targets = inputs [ i ] , targets [ i ] preds = apply_fun ( params , rng , inputs ) . squeeze ( ) mean_loss = ( - norm . logpdf ( targets . squeeze ( ) , preds , params [ "noise" ] ) . mean ( ) if loss_type == "nlp" else np . power ( targets . squeeze ( ) - preds , 2 ) . mean ( ) ) if reduce == "sum" : loss = n * mean_loss elif reduce == "mean" : loss = mean_loss return loss def sample_fun_fun ( rng , params ) : def f ( x ) : return apply_fun ( params , rng , x ) return f return { "init" : init_fun , "apply" : apply_fun , "gaussian" : gaussian_fun , "loss" : loss_fun , "sample_function" : sample_fun_fun , }
short_docstring: Creates init and apply functions for a fully-connected NN with a Gaussian
  likelihood.
short_docstring_tokens: Creates init and apply functions for a fully - connected NN with a Gaussian likelihood .
comment: """
      Creates init and apply functions for a fully-connected NN with a Gaussian
      likelihood.
  
      :param output_dim: Number of output dimensions
      :param first_layer_factor: Scaling factor for first dense layer initialization.
      """$SEP$# No skips on the last one!$SEP$# Conform to API:$SEP$"""
          :param batch_size: How large a batch to subselect from the provided data
          :param n: The total size of the dataset (to multiply batch estimate by)
          """