MixtureTable:updateOutput

1 files changed, 39 insertions, 0 deletions

diff --git a/MixtureTable.lua b/MixtureTable.lua
new file mode 100644
index 0000000..5403800
--- /dev/null
+++ b/MixtureTable.lua
@@ -0,0 +1,39 @@
+local MixtureTable, parent = torch.class('nn.MixtureTable', 'nn.Module')
+
+function MixtureTable:__init()
+   parent.__init(self)
+   self._gate = torch.Tensor()
+   self.size = torch.LongTensor()
+   self.batchSize = 0
+end
+
+function MixtureTable:updateOutput(input) 
+   local gaterInput, expertInputs = unpack(input)
+   if gaterInput:dim() == 2 then
+      if gaterInput:size(2) ~= #expertInputs then
+         error"Should be one gater output per expert"
+      end
+      local expertInput = expertInputs[1]
+      if self.batchSize ~= expertInput:size(1) then
+         self.size:resize(expertInput:dim()):fill(1)
+         self.size[1] = expertInput:size(1)
+         self.output:resizeAs(expertInput)
+         self.batchSize ~= expertInput:size(1)
+      end
+      self.output:zero()
+      for i,expertInput in ipairs(expertInputs) do
+         -- multiply each gater output (a gate) by its 
+         -- commensurate expert
+         self._gate:resize(self.size:storage())
+         self._gate:copy(gaterInput:select(2,i))
+         self.output:addcmul(expertInput,gate:expandAs(expertInput))
+      end
+   end
+   
+   return self.output
+end
+
+function MixtureTable:updateGradInput(input, gradOutput)
+   
+   return self.gradInput
+end