docstr: :param input_memories: [num-inputs-flattened, D]
  :param input_memories_origin_idx: [num-inputs-flattened]
  :param initial_states: [num-targets, H]
  :param target_token_ids: [num-targets, max-seq-size]
  :param target_lengths: [num-targets]
  :param copyable_elements_idxs: [num-copyable-elements]
  :param copyable_elements_sample_idxs: [num-copyable-elements]
  
  :return: the loss function.
docstr_tok: [":","param","input_memories",":","[","num-inputs-flattened",",","D","]",":","param","input_memories_origin_idx",":","[","num-inputs-flattened","]",":","param","initial_states",":","[","num-targets",",","H","]",":","param","target_token_ids",":","[","num-targets",",","max-seq-size","]",":","param","target_lengths",":","[","num-targets","]",":","param","copyable_elements_idxs",":","[","num-copyable-elements","]",":","param","copyable_elements_sample_idxs",":","[","num-copyable-elements","]",":","return",":","the","loss","function","."]
code: def forward(
      self,
      *,
      input_memories,
      input_memories_origin_idx,
      initial_states,
      target_token_ids,
      copyable_elements_idxs,
      copyable_elements_sample_idxs,
      target_lengths,
  ):
      """
      :param input_memories: [num-inputs-flattened, D]
      :param input_memories_origin_idx: [num-inputs-flattened]
      :param initial_states: [num-targets, H]
      :param target_token_ids: [num-targets, max-seq-size]
      :param target_lengths: [num-targets]
      :param copyable_elements_idxs: [num-copyable-elements]
      :param copyable_elements_sample_idxs: [num-copyable-elements]
  
      :return: the loss function.
      """
      copy_logprobs, target_logprobs, _ = self._compute_logprobs(
          initial_states,
          input_memories,
          input_memories_origin_idx,
          target_token_ids[:, :-1],
      )
  
      # Get loss. UNKs are only predicted if we cannot copy.
      num_valid_copy_actions = scatter_add(
          src=torch.ones_like(copyable_elements_sample_idxs),
          index=copyable_elements_sample_idxs,
          dim_size=target_token_ids.shape[0]
          * (target_token_ids.shape[1] - 1),
      )
      locations_with_valid_copy_actions = (
          num_valid_copy_actions.reshape(
              target_token_ids.shape[0], target_token_ids.shape[1] - 1
          )
          > 0
      )
      unk_prediction_locations = (
          target_token_ids[:, 1:] == self.__unk_id
      )
      mask = (
          locations_with_valid_copy_actions & unk_prediction_locations
      )
  
      correct_generation_logprobs = torch.gather(
          target_logprobs,
          index=target_token_ids[:, 1:].unsqueeze(-1),
          dim=-1,
      ).squeeze(
          -1
      )  # [num-targets, max-seq-size -1]
      correct_generation_logprobs.masked_fill_(mask, -math.inf)
  
      correct_copy_logprobs = scatter_logsumexp(
          src=copy_logprobs.flatten()[copyable_elements_idxs],
          index=copyable_elements_sample_idxs,
          dim=0,
          dim_size=target_token_ids.shape[0]
          * (target_token_ids.shape[1] - 1),
          eps=0,
      )  # [num-targets * (max-seq-size-1)]
      correct_copy_logprobs = correct_copy_logprobs.view(
          target_token_ids.shape[0], target_token_ids.shape[1] - 1
      )
  
      any_correct_action_logprob = torch.logsumexp(
          torch.stack(
              (correct_generation_logprobs, correct_copy_logprobs)
          ),
          dim=0,
      )  # [num-targets, max-seq-size]
  
      mask = torch.arange(
          any_correct_action_logprob.shape[1],
          device=target_lengths.device,
      ).unsqueeze(0) < target_lengths.unsqueeze(1)
      per_seq_loss = (any_correct_action_logprob * mask.float()).sum(
          dim=-1
      ) / mask.float().sum(dim=-1)
  
      return -per_seq_loss.mean()
  
imports: from typing_extensions import Final
  import logging
  import math
  import numpy
  import torch
  import torch.nn
  from dpu_utils.mlutils import Vocabulary
  from torch_scatter import scatter_add
  from torch_scatter.composite import scatter_log_softmax, scatter_logsumexp
  from typing import Any, Counter, Dict, List, NamedTuple, Tuple, Union
  from ptgnn.baseneuralmodel import AbstractNeuralModel
label_desc: Use image node_id 3 for calling the GruCopyingDecoder obj's underlying member method code with example usage: obj.forward() and returns: unknown