code: def __init__(
      self,
      model: nn.Module,
      shape: Union[None, Tuple[int, int], List[int]] = (16, 16),
      regularization_key: str = "regularization",
      encoder_kwargs: Optional[Dict[str, Any]] = None,
  ):
      super().__init__()
      self.model = model
      assert hasattr(model, "encode") and hasattr(
          model, "decode"
      ), "Need AE interface"
      self.regularization = get_nested_attribute(
          model, regularization_key
      )
      self.shape = shape
      self.encoder_kwargs = default(
          encoder_kwargs, {"return_reg_log": True}
      )
  
imports: import logging
  import math
  import re
  from abc import abstractmethod
  from contextlib import contextmanager
  from typing import Any, Dict, List, Optional, Tuple, Union
  import pytorch_lightning
  import torch
  import torch.nn
  from einops import rearrange
  from packaging import version
  from ..modules.autoencoding.regularizers import AbstractRegularizer
  from ..modules.ema import LitEma
  from ..util import default, get_nested_attribute, get_obj_from_str, instantiate_from_config
label_desc: Use image node_id 1 to create a new AEIntegerWrapper object from inherited base classes: nn with example: obj = AEIntegerWrapper(model, shape, regularization_key, encoder_kwargs)