class_docstr: Noisy Network from Fortunato et al. (2018).
  
  Attributes:
    rng_key: jax.interpreters.xla.DeviceArray, key for JAX RNG.
    eval_mode: bool, whether to turn off noise during evaluation.
code: def __call__(self, x, features, bias=True, kernel_init=None):
      def mu_init(key, shape):
          # Initialization of mean noise parameters (Section 3.2)
          low = -1 / jnp.power(x.shape[0], 0.5)
          high = 1 / jnp.power(x.shape[0], 0.5)
          return jax.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.1 / onp.sqrt(x.shape[0]))
  
      if self.eval_mode:
          # Turn off noise during evaluation
          w_epsilon = onp.zeros(
              shape=(x.shape[0], features), dtype=onp.float32
          )
          b_epsilon = onp.zeros(shape=(features,), dtype=onp.float32)
      else:
          # Factored gaussian noise in (10) and (11) in Fortunato et al. (2018).
          rng_p, rng_q = jax.random.split(self.rng_key, num=2)
          p = NoisyNetwork.sample_noise(rng_p, [x.shape[0], 1])
          q = NoisyNetwork.sample_noise(rng_q, [1, 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", mu_init, (x.shape[0], features))
      w_sigma = self.param(
          "kernel_sigma", sigma_init, (x.shape[0], features)
      )
      w = w_mu + jnp.multiply(w_sigma, w_epsilon)
      ret = jnp.matmul(x, w)
  
      b_mu = self.param("bias_mu", mu_init, (features,))
      b_sigma = self.param("bias_sigma", sigma_init, (features,))
      b = b_mu + jnp.multiply(b_sigma, b_epsilon)
      return jnp.where(bias, ret + b, ret)
  
imports: import itertools
  import time
  from typing import Optional, Sequence, Tuple, Union
  from absl import logging
  from dopamine.discrete_domains import atari_lib
  from flax import linen
  import gin
  import jax
  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, features, bias, kernel_init) and returns: jnp, jax, unknown