class_docstr: SpectralNormalizedLinear module.
  
  This is a Linear layer with a upper-bounded Lipschitz. It is used in iResNet.
  
  Reference:
    Behrmann et al. Invertible Residual Networks. ICML 2019.
    https://arxiv.org/pdf/1811.00995.pdf
class_docstr_tok: ["SpectralNormalizedLinear","module",".","This","is","a","Linear","layer","with","a","upper-bounded","Lipschitz",".","It","is","used","in","iResNet",".","Reference",":","Behrmann","et","al",".","Invertible","Residual","Networks",".","ICML","2019.","https",":","\/\/arxiv.org\/pdf\/1811.00995.pdf"]
code: def get_normalized_weights(
      self, weights: jnp.ndarray, renormalize: bool = False
  ) -> jnp.ndarray:
      def _l2_normalize(x, axis=None, eps=1e-12):
          return x * jax.lax.rsqrt(
              (x * x).sum(axis=axis, keepdims=True) + eps
          )
  
      output_size = self.output_size
      dtype = weights.dtype
      assert output_size == weights.shape[-1]
      sigma = hk.get_state("sigma", (), init=jnp.ones)
      if renormalize:
          # Power iterations to compute spectral norm V*W*U^T.
          u = hk.get_state(
              "u",
              (1, output_size),
              dtype,
              init=hk.initializers.RandomNormal(),
          )
          for _ in range(self.num_iterations):
              v = _l2_normalize(
                  jnp.matmul(u, weights.transpose()), eps=self.eps
              )
              u = _l2_normalize(jnp.matmul(v, weights), eps=self.eps)
          u = jax.lax.stop_gradient(u)
          v = jax.lax.stop_gradient(v)
          sigma = jnp.matmul(jnp.matmul(v, weights), jnp.transpose(u))[
              0, 0
          ]
          hk.set_state("u", u)
          hk.set_state("v", v)
          hk.set_state("sigma", sigma)
      factor = jnp.maximum(1, sigma / self.lipschitz_coeff)
      return weights / factor
  
imports: import dataclasses
  import functools
  from typing import Any, Callable, Generic, Iterable, Optional
  from acme import specs
  from acme import types
  from acme.jax import networks
  from acme.jax import utils
  from acme.jax.imitation_learning_types import DirectRLNetworks
  import haiku
  import jax
  from jax import numpy
  import numpy
label_desc: Use image node_id 2 for calling the SpectralNormalizedLinear obj's underlying member method code with example usage: obj.get_normalized_weights(weights, renormalize) and returns: unknown, unknown