class_docstr: MuZero-style transition model cell, used for SPR.
  
  Attributes:
    num_actions: how many actions are possible (shape of one-hot vector)
    latent_dim: number of channels in representation.
    renormalize: whether or not to apply renormalization.
code: def setup(self):
      self.bn = NoStatsBatchNorm(axis=-1, axis_name="batch")
  
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 1 for calling the ConvTMCell obj's underlying member method code with example usage: obj.setup() without return types