<commit_before>local RIndexLinear, parent = torch.class('onmt.RIndexLinear', 'nn.Linear')

--[[

  MaskLinear is a Linear layer that allows the user to a set a collection of
  output row indices. When the row indices are set, the layer will behave like a
  Linear layer that output only these rows - all the other ones will be zero.
  This is particularily interesting for softmax reduction.

]]--

function RIndexLinear:__init(inputSize, outputSize, bias)
  parent.__init(self, inputSize, outputSize, bias)
  self.restrictedOutput = torch.Tensor(1)
end

function RIndexLinear:setOutputIndices(rowIndices)
  if self.rowIndices ~= rowIndices then
    self.rowIndices = rowIndices
    if rowIndices then
      self.restrictedWeight = self.weight:index(1, rowIndices)
      self.restrictedBias = self.bias:index(1, rowIndices)
    end
  end
end

local function updateAddBuffer(self, input)
   local nbatch = input:size(1)
   self.addBuffer = self.addBuffer or input.new()
   if self.addBuffer:nElement() ~= nbatch then
      self.addBuffer:resize(nbatch):fill(1)
   end
end

function RIndexLinear:updateOutput(input)
  if self.train or not self.rowIndices then
    return parent.updateOutput(self, input)
  end
  if input:dim() == 1 then
    self.output:resize(self.weight:size(1)):zero()
    self.restrictedOutput:resize(self.restrictedWeight:size(1))
    if self.bias then
      self.restrictedOutput:copy(self.restrictedBias)
    else
      self.restrictedOutput:zero()
    end
    self.restrictedOutput:addmv(1, self.restrictedWeight, input)
    self.output:indexCopy(1, self.rowIndices, self.restrictedOutput)
  elseif input:dim() == 2 then
    local nbatch = input:size(1)<commit_msg>invalid trigger for RIndexLinear<commit_after>local RIndexLinear, parent = torch.class('onmt.RIndexLinear', 'nn.Linear')

--[[

  MaskLinear is a Linear layer that allows the user to a set a collection of
  output row indices. When the row indices are set, the layer will behave like a
  Linear layer that output only these rows - all the other ones will be zero.
  This is particularily interesting for softmax reduction.

]]--

function RIndexLinear:__init(inputSize, outputSize, bias)
  parent.__init(self, inputSize, outputSize, bias)
  self.restrictedOutput = torch.Tensor(1)
end

function RIndexLinear:setOutputIndices(rowIndices)
  if self.rowIndices ~= rowIndices then
    self.rowIndices = rowIndices
    if rowIndices then
      self.restrictedWeight = self.weight:index(1, rowIndices)
      self.restrictedBias = self.bias:index(1, rowIndices)
    end
  end
end

local function updateAddBuffer(self, input)
   local nbatch = input:size(1)
   self.addBuffer = self.addBuffer or input.new()
   if self.addBuffer:nElement() ~= nbatch then
      self.addBuffer:resize(nbatch):fill(1)
   end
end

function RIndexLinear:updateOutput(input)
  if not self.train or not self.rowIndices then
    return parent.updateOutput(self, input)
  end
  if input:dim() == 1 then
    self.output:resize(self.weight:size(1)):zero()
    self.restrictedOutput:resize(self.restrictedWeight:size(1))
    if self.bias then
      self.restrictedOutput:copy(self.restrictedBias)
    else
      self.restrictedOutput:zero()
    end
    self.restrictedOutput:addmv(1, self.restrictedWeight, input)
    self.output:indexCopy(1, self.rowIndices, self.restrictedOutput)
  elseif input:dim() == 2 then
    local nbatch = input:size(1)
{'fraction_non_alphanumeric': 0.0936120105975861, 'fraction_numerical': 0.007065057403591404, 'mean_word_length': 4.284603421461897, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3686764', 'n_tokens_mistral': 1063, 'n_tokens_neox': 1035, 'n_words': 289}