class_docstr: A Jax implementation of the SPR transition model, leveraging scan.
  
  Attributes:
    num_actions: How many possible actions exist.
    latent_dim: Output size.
    renormalize: Whether or not to apply renormalization.
code: def __call__(self, x, action):
      scan = nn.scan(
          ConvTMCell,
          in_axes=0,
          out_axes=0,
          variable_broadcast=["params"],
          split_rngs={"params": False},
      )(
          latent_dim=self.latent_dim,
          num_actions=self.num_actions,
          renormalize=self.renormalize,
      )
      return scan(x, action)
  
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 TransitionModel obj's underlying member method code with example usage: obj.__call__(x, action) and returns: scan