class_docstr: Noisy Network from Fortunato et al. (2018).
docstr: Call the noisy layer.
  
  Args:
    x: Data point. jnp.float32 tensor, without batch dimension
    rng_key: JAX prng key
    bias: Whether or not to use bias params (static)
    kernel_init: Init function for kernels
    eval_mode: Enable eval mode. Disables noise parameters.
  
  Returns:
    The transformed output. JNP tensor.
docstr_tok: ["Call","the","noisy","layer",".","Args",":","x",":","Data","point",".","jnp.float32","tensor",",","without","batch","dimension","rng_key",":","JAX","prng","key","bias",":","Whether","or","not","to","use","bias","params","(","static",")","kernel_init",":","Init","function","for","kernels","eval_mode",":","Enable","eval","mode",".","Disables","noise","parameters",".","Returns",":","The","transformed","output",".","JNP","tensor","."]
code: def __call__(
      self, x, rng_key, bias=True, kernel_init=None, eval_mode=False
  ):
      """Call the noisy layer.
  
      Args:
        x: Data point. jnp.float32 tensor, without batch dimension
        rng_key: JAX prng key
        bias: Whether or not to use bias params (static)
        kernel_init: Init function for kernels
        eval_mode: Enable eval mode. Disables noise parameters.
  
      Returns:
        The transformed output. JNP tensor.
      """
  
      def mu_init(key, shape):
          # Initialization of mean noise parameters (Section 3.2)
          low = -1 / jnp.power(x.shape[-1], 0.5)
          high = 1 / jnp.power(x.shape[-1], 0.5)
          return random.uniform(
              key, minval=low, maxval=high, shape=shape
          )
  
      def sigma_init(
          key, shape, dtype=jnp.float32
      ):  # pylint: disable=unused-argument
          # Initialization of sigma noise parameters (Section 3.2)
          return jnp.ones(shape, dtype) * (0.5 / onp.sqrt(x.shape[-1]))
  
      # Factored gaussian noise in (10) and (11) in Fortunato et al. (2018).
      p = NoisyNetwork.sample_noise(rng_key, [x.shape[-1], 1])
      q = NoisyNetwork.sample_noise(rng_key, [1, self.features])
      f_p = NoisyNetwork.f(p)
      f_q = NoisyNetwork.f(q)
      w_epsilon = f_p * f_q
      b_epsilon = jnp.squeeze(f_q)
  
      # See (8) and (9) in Fortunato et al. (2018) for output computation.
      w_mu = self.param("kernel", mu_init, (x.shape[-1], self.features))
      w_sigma = self.param(
          "kernell", sigma_init, (x.shape[-1], self.features)
      )
      w_epsilon = jnp.where(
          eval_mode,
          onp.zeros(
              shape=(x.shape[-1], self.features), dtype=onp.float32
          ),
          w_epsilon,
      )
      w = w_mu + jnp.multiply(w_sigma, w_epsilon)
      ret = jnp.matmul(x, w)
  
      b_epsilon = jnp.where(
          eval_mode,
          onp.zeros(shape=(self.features,), dtype=onp.float32),
          b_epsilon,
      )
      b_mu = self.param("bias", mu_init, (self.features,))
      b_sigma = self.param("biass", sigma_init, (self.features,))
      b = b_mu + jnp.multiply(b_sigma, b_epsilon)
      return jnp.where(bias, ret + b, ret)
  
imports: import collections
  import functools
  import time
  from typing import Any, Callable, Optional, Tuple
  from dopamine.jax.networks import preprocess_atari_inputs
  from flax import linen
  import gin
  import jax
  from jax import lax
  from jax import random
  from jax.nn import initializers
  import jax.numpy
  import numpy
label_desc: Use image node_id 3 for calling the NoisyNetwork obj's underlying member method code with example usage: obj.__call__(x, rng_key, bias, kernel_init, eval_mode) and returns: jnp, random, unknown