class_docstr: FID model with a retrieval module over a knowledge memory.
class_docstr_tok: ["FID","model","with","a","retrieval","module","over","a","knowledge","memory","."]
docstr: Returns negative loglikelihood (NLL) of the target sentence.
  
  Args:
    output: Output of model in OrderedDict.
    batch: Batch of data that has 'decoder_target' as ground-truth.
  
  Returns:
    Total loss.
docstr_tok: ["Returns","negative","loglikelihood","(","NLL",")","of","the","target","sentence",".","Args",":","output",":","Output","of","model","in","OrderedDict",".","batch",":","Batch","of","data","that","has","'decoder_target","'","as","ground-truth",".","Returns",":","Total","loss","."]
code: def loss_function_dict(
      self, output: constants.JTensorDict, batch: constants.JTensorDict
  ) -> Dict[str, Any]:
      """Returns negative loglikelihood (NLL) of the target sentence.
  
      Args:
        output: Output of model in OrderedDict.
        batch: Batch of data that has 'decoder_target' as ground-truth.
  
      Returns:
        Total loss.
      """
      model_config = self.config.model
      gen_loss = losses.nll_loss(
          targets=batch["decoder_target_tokens"],
          pred=output["predicted_logits"],
          target_masks=batch["decoder_target_tokens"] > 0,
          label_smoothing=self.config.model.get("label_smoothing"),
      )
      loss_dict = {"gen_loss": gen_loss}
  
      if "inbatch_sim" in output:
          score_matrix = output["inbatch_sim"]
          bsz = score_matrix.shape[0]
          labels = jnp.arange(bsz) + bsz * jax.lax.axis_index(
              axis_name="batch"
          )
          contra_loss = losses.nll_loss(
              pred=score_matrix / self.config.model.get("temperature"),
              targets=labels,
          )
          loss_dict["contra_loss"] = contra_loss
          r = model_config.retrieval_ratio
          loss = gen_loss * (1 - r) + contra_loss * r
          accs = jnp.equal(jnp.argmax(score_matrix, axis=1), labels)
          loss_dict["contra_accs"] = accs
      else:
          loss_dict["contra_loss"] = 0.0
          loss_dict["contra_accs"] = 0.0
          loss = gen_loss
  
      if "disentangle" in model_config and "disentangle_reg" in output:
          loss += output["disentangle_reg"] * 1e-2
      if "gap" in model_config and "gap" in output:
          loss += output["gap"] * 1e-4
      loss_dict["total_loss"] = loss
      return loss_dict
  
imports: import functools
  from typing import Any, Dict, Mapping, Optional, Tuple, List
  from absl import logging
  import flax.linen
  import jax
  from jax.experimental import host_callback
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.projects.knowledge_visual_language.data import data_utils
  from scenic.projects.knowledge_visual_language.models import constants
  from scenic.projects.knowledge_visual_language.models import fusion_in_decoder_soft
  from scenic.projects.knowledge_visual_language.models import layers
  from scenic.projects.knowledge_visual_language.models import local_memory
  from scenic.projects.knowledge_visual_language.models import losses
  from scenic.projects.knowledge_visual_language.models import metrics
  from t5x import decoding
label_desc: Use image node_id 3 for calling the KnowledgeFIDModel obj's underlying member method code with example usage: obj.loss_function_dict(output, batch) and returns: loss_dict