parameters: (
          self, init_fun: Optional[NNInitFunc] = None, *, seed: Optional[PRNGKeyT] = None
      )
docstring: r"""
          Re-initializes all the parameters with the provided initialization function,
          defaulting to the normal distribution of standard deviation 0.01.
  
          .. warning::
  
              The init function will not change the dtype of the parameters, which is
              determined by the model. DO NOT SPECIFY IT INSIDE THE INIT FUNCTION
  
          Args:
              init_fun: a jax initializer such as :ref:`jax.nn.initializers.normal`.
                  Must be a Callable taking 3 inputs, the jax PRNG key, the shape and the
                  dtype, and outputting an array with the valid dtype and shape. If left
                  unspecified, defaults to :code:`jax.nn.initializers.normal(stddev=0.01)`
              seed: Optional seed to be used. The seed is synced across all MPI processes.
                  If unspecified, uses a random seed.
docstring_summary: r"""
          Re-initializes all the parameters with the provided initialization function,
          defaulting to the normal distribution of standard deviation 0.01.
docstring_tokens: r Re - initializes all the parameters with the provided initialization function defaulting to the normal distribution of standard deviation 0 . 01 .
function: def init_parameters(
          self, init_fun: Optional[NNInitFunc] = None, *, seed: Optional[PRNGKeyT] = None
      ):
          r"""
          Re-initializes all the parameters with the provided initialization function,
          defaulting to the normal distribution of standard deviation 0.01.
  
          .. warning::
  
              The init function will not change the dtype of the parameters, which is
              determined by the model. DO NOT SPECIFY IT INSIDE THE INIT FUNCTION
  
          Args:
              init_fun: a jax initializer such as :ref:`jax.nn.initializers.normal`.
                  Must be a Callable taking 3 inputs, the jax PRNG key, the shape and the
                  dtype, and outputting an array with the valid dtype and shape. If left
                  unspecified, defaults to :code:`jax.nn.initializers.normal(stddev=0.01)`
              seed: Optional seed to be used. The seed is synced across all MPI processes.
                  If unspecified, uses a random seed.
          """
          if init_fun is None:
              init_fun = normal(stddev=0.01)
  
          rng = nkjax.PRNGSeq(nkjax.PRNGKey(seed))
  
          def new_pars(par):
              return jnp.asarray(
                  init_fun(rng.take(1)[0], shape=par.shape, dtype=par.dtype),
                  dtype=par.dtype,
              )
  
          self.parameters = jax.tree_map(new_pars, self.parameters)
function_tokens: def init_parameters ( self , init_fun : Optional [ NNInitFunc ] = None , * , seed : Optional [ PRNGKeyT ] = None ) : if init_fun is None : init_fun = normal ( stddev = 0.01 ) rng = nkjax . PRNGSeq ( nkjax . PRNGKey ( seed ) ) def new_pars ( par ) : return jnp . asarray ( init_fun ( rng . take ( 1 ) [ 0 ] , shape = par . shape , dtype = par . dtype ) , dtype = par . dtype , ) self . parameters = jax . tree_map ( new_pars , self . parameters )