-- A module that takes in an array of tensors, allows each component module to potentially receive all inputs, and dynamically composes an output array of tensors by concatenating a single tensor output from each component module.

-- Each module added to a PDT takes in the table of all inputs and returns a single tensor as output. Each module is thus expected to produce a table of gradInputs matching the table of inputs. However, the entries of this table are allowed to be nil, in which case they are treated like a tensor of all zeros. In general, PDT should only be used in conjunction with SelectTables, which route both the inputs and gradInputs appropriately. We'd like to avoid allocating new memory each time updateGradInput is called, so gradInputs will maintain its own gradInput tensors when necessary, and add from the component modules. If a given input only feeds into a single module, the modules gradOutput is passed through directly, rather than copied to local storage, for efficiency.

-- It's not possible to split the output of a module within a ParallelTable, since the gradOutputs don't align with the modules, and cannot be routed one-to-one

local ParallelDistributingTable, parent = torch.class('nn.ParallelDistributingTable', 'nn.ParallelTable')

function ParallelDistributingTable:__init(name)

parent.__init(self)

self.num_grad_input_sources = torch.Tensor() -- keeps track of how many modules are connected to a given input, so gradOutput from a module can be routed directly to gradInput if only one module connects to the corresponding input

self.name = name

end

function ParallelDistributingTable:updateOutput(input)

for i=1,#self.modules do

-- it would be easy and no less efficient to append a table of outputs from the current module to the end of the list. However, we would then need to keep track of which gradOutputs get routed to which module

self.output[i] = self.modules[i]:updateOutput(input) -- rather than input[i], in ParallelTable

end

return self.output

end

function ParallelDistributingTable:updateGradInput(input, gradOutput)

-- determine the number of modules connected to each entry of the input table. If an input tensor only feeds into one module, the corresponding gradOutput can be copied directly into the PDT's gradInput. Otherwise, we need to maintain a distinct gradInput tensor and add the contribution from each connected module.

self.num_grad_input_sources:resize(#input)

self.num_grad_input_sources:zero()

for i,module in ipairs(self.modules) do

for j in pairs(module:updateGradInput(input, gradOutput[i])) do

self.num_grad_input_sources[j] = self.num_grad_input_sources[j] + 1

end

end

-- If we're not routing directly from gradOutput, resize the gradInputs based on the inputs, and zero.

for i,input_tensor_i in ipairs(input) do

if self.num_grad_input_sources[i] == 1 then

self.gradInput[i] = nil -- if only one module connects to input i, pass the gradOutput through directly

else

if not(self.gradInput[i]) then

self.gradInput[i] = input_tensor_i.new() -- create a new tensor of the same type as the input

end

self.gradInput[i]:resizeAs(input_tensor_i):zero()

end

end

-- Iterate over the modules. Within each module, iterate over the entries of gradInputs using pairs (since some entries will be nil). For all existent entries, add the gradInput into the PDT's gradInput

for i,module in ipairs(self.modules) do

for j,gradInput_tensor_j in pairs(module.gradInput) do -- updateGradInput was already called before

if self.gradInput[j] == nil then

self.gradInput[j] = gradInput_tensor_j

else

self.gradInput[j]:add(gradInput_tensor_j)

end

end

end

return self.gradInput

end

function ParallelDistributingTable:accGradParameters(input, gradOutput, scale)

scale = scale or 1

for i,module in ipairs(self.modules) do

module:accGradParameters(input, gradOutput[i], scale) -- rather than input[i]

end

end

function ParallelDistributingTable:accUpdateGradParameters(input, gradOutput, lr)

lr = lr or 1

for i,module in ipairs(self.modules) do

module:accUpdateGradParameters(input, gradOutput[i], lr) -- rather than input[i]

end

end

function ParallelDistributingTable:__tostring__()

local tab = ' '

local line = '\n'

local next = ' |`-> '

local ext = ' | '

local extlast = ' '

local last = ' ... -> '

local str = 'nn.ParallelDistributingTable'

if self.name then

str = str .. ' (' .. self.name .. ')'

end

str = str .. ' {' .. line .. tab .. 'input'

for i=1,#self.modules do

if i == #self.modules then

str = str .. line .. tab .. next .. '(' .. i .. '): ' .. tostring(self.modules[i]):gsub(line, line .. tab .. extlast)

else

str = str .. line .. tab .. next .. '(' .. i .. '): ' .. tostring(self.modules[i]):gsub(line, line .. tab .. ext)

end

end

str = str .. line .. tab .. last .. 'output'

str = str .. line .. '}'

return str

end