class_docstr: SockeyeModel shares components needed for both training and inference.
  The main components of a Sockeye model are
  1) Source embedding
  2) Target embedding
  3) Encoder
  4) Decoder
  5) Output Layer
  
  ModelConfig contains parameters and their values that are fixed at training time and must be re-used at inference
  time.
  
  :param config: Model configuration.
  :param inference_only: Use the model only for inference, enabling optimizations.
  :param clamp_to_dtype: Avoid -inf/inf by clamping outputs to min/max finite values for their dtype.
class_docstr_tok: ["SockeyeModel","shares","components","needed","for","both","training","and","inference",".","The","main","components","of","a","Sockeye","model","are","1",")","Source","embedding","2",")","Target","embedding","3",")","Encoder","4",")","Decoder","5",")","Output","Layer","ModelConfig","contains","parameters","and","their","values","that","are","fixed","at","training","time","and","must","be","re-used","at","inference","time",".",":","param","config",":","Model","configuration",".",":","param","inference_only",":","Use","the","model","only","for","inference",",","enabling","optimizations",".",":","param","clamp_to_dtype",":","Avoid","-inf\/inf","by","clamping","outputs","to","min\/max","finite","values","for","their","dtype","."]
code: def __init__(
      self,
      config: ModelConfig,
      inference_only: bool = False,
      clamp_to_dtype: bool = False,
      train_decoder_only: bool = False,
      forward_pass_cache_size: int = 0,
  ) -> None:
      super().__init__()
      self.config = copy.deepcopy(config)
      self.dtype = utils.get_torch_dtype(config.dtype)
      self.clamp_to_dtype = clamp_to_dtype
      logger.info("%s", self.config)
      self.train_decoder_only = train_decoder_only
      self.forward_pass_cache_size = forward_pass_cache_size
      self.embed_and_encode = self._embed_and_encode
      if self.forward_pass_cache_size > 0:
          self.embed_and_encode = self._cache_wrapper(
              self._embed_and_encode
          )
  
      # source & target embeddings, potentially shared/tied
      (
          source_embedding,
          target_embedding,
          output_weight,
      ) = self._get_embeddings()
  
      self.embedding_source = encoder.Embedding(
          config.config_embed_source,
          embedding=source_embedding,
          dtype=self.dtype,
      )
      self.embedding_target = encoder.Embedding(
          config.config_embed_target,
          embedding=target_embedding,
          dtype=self.dtype,
      )
  
      # encoder & decoder first (to know the decoder depth)
      self.encoder = encoder.get_transformer_encoder(
          self.config.config_encoder,
          inference_only=inference_only,
          dtype=self.dtype,
          clamp_to_dtype=clamp_to_dtype,
      )
      self.decoder = decoder.get_decoder(
          self.config.config_decoder,
          inference_only=inference_only,
          dtype=self.dtype,
          clamp_to_dtype=clamp_to_dtype,
      )
      self.nvs = None
      if self.config.neural_vocab_selection:
          self.nvs = nvs.NeuralVocabSelection(
              model_size=self.config.config_encoder.model_size,
              vocab_target_size=self.config.vocab_target_size,
              model_type=self.config.neural_vocab_selection,
              dtype=self.dtype,
          )
  
      self.output_layer = layers.OutputLayer(
          hidden_size=self.decoder.get_num_hidden(),
          vocab_size=self.config.vocab_target_size,
          weight=output_weight,
          dtype=self.dtype,
      )
      self.output_layer_module_cached = self.output_layer
      # Running this layer scripted with a newly initialized model can cause an overflow error.
      self.output_layer_script_cached = pt.jit.script(
          self.output_layer_module_cached
      )
      self.set_inference_only(inference_only)
  
      self.factor_output_layers = pt.nn.ModuleList()
      # Optional target factor output layers
      for i, factor_config in enumerate(self.target_factor_configs, 1):
          # Each target stream has its own, independent output layer
          # TODO also consider weight tying with target factor input embeddings
          output_layer = pt.nn.Linear(
              in_features=self.decoder.get_num_hidden(),
              out_features=factor_config.vocab_size,
              bias=True,
              dtype=self.dtype,
          )
          self.factor_output_layers.append(output_layer)
      self.factor_vocab_size = (
          factor_config.vocab_size
          if self.target_factor_configs
          else None
      )
  
      self.length_ratio = None  # type: Optional[layers.LengthRatio]
      if self.config.config_length_task is not None:
          utils.check_condition(
              self.config.config_length_task.weight > 0.0,
              "Auxiliary length task requested, but its loss weight is zero",
          )
          self.length_ratio = layers.LengthRatio(
              hidden_size=self.encoder.get_num_hidden(),
              num_layers=self.config.config_length_task.num_layers,
              dtype=self.dtype,
          )
  
      # traced components (for inference)
      self.traced_embedding_source = (
          None
      )  # type: Optional[pt.jit.ScriptModule]
      self.traced_encoder = None  # type: Optional[pt.jit.ScriptModule]
      self.traced_decode_step = (
          None
      )  # type: Optional[pt.jit.ScriptModule]
  
      # Make sure all parameters are in the model's specified dtype. Warn when
      # conversion is required. This is a no-op when all submodules are
      # instantiated using the model's dtype.
      mismatched_dtype_params = [
          (name, param.dtype)
          for name, param in self.named_parameters()
          if param.dtype != self.dtype
      ]
      self.to(self.dtype)
      if mismatched_dtype_params:
          logger.warn(
              "Some parameters were created in a different dtype and then converted to the SockeyeModel's "
              "dtype. This works but can cause memory spikes when creating/loading models. To avoid this, "
              "pass the SockeyeModel's dtype when instantiating all submodules. Converted parameters:"
          )
          for name, dtype in mismatched_dtype_params:
              logger.warn(f"{name}: {dtype} -> {self.dtype}")
  
      self.knn: Optional[layers.KNN] = None
  
imports: import copy
  import logging
  import os
  import time
  from dataclasses import dataclass
  from functools import lru_cache
  from typing import cast, Dict, List, Optional, Tuple, Union
  import numpy
  import torch
  from sockeye import __version__
  from .None import constants
  from .None import data_io
  from .None import decoder
  from .None import encoder
  from .None import layers
  from .None import transformer
  from .None import utils
  from .None import vocab
  from .config import Config
  from .encoder import FactorConfig
  from .layers import LengthRatioConfig
  from .None import nvs
  from sockeye.knn import KNNConfig
label_desc: Use image node_id 1 to create a new SockeyeModel object from inherited base classes: pt.nn with example: obj = SockeyeModel(config, inference_only, clamp_to_dtype, train_decoder_only, forward_pass_cache_size)