Give LOGO decoders the ability to track state across calls.

src/main/scala/epic/logo/ArgmaxInferencer.scala

@@ -1,7 +1,9 @@
 package epic.logo
 
-trait ArgmaxInferencer[T, Y, W] {
-
-  def argmax(weights : Weights[W], instance : T) : (Y, W, Double)
+trait ArgmaxInferencer[T, W, S] extends Inferencer[S] {
+  type Y
 
+  def argmax(weights : Weights[W], instance : T) : (Y, W, Double, S)
+  def initialState: S
+  def reduceStates(s1: S, s2: S): S
 }

src/main/scala/epic/logo/CompoundIterationCallback.scala

@@ -1,27 +1,31 @@
 package epic.logo
 
-class CompoundIterationCallback[T, W](val callbacks : Iterable[IterationCallback[T, W]]) extends IterationCallback[T, W] {
+class CompoundIterationCallback[T, W, S1, S2](val callbacks: Iterable[IterationCallback[T, W, S1, S2]])
+    extends IterationCallback[T, W, S1, S2] {
 
-  override def startIteration(iter : Int, weights : Weights[W]) : Unit = {
+  override def startIteration(iter: Int, weights: Weights[W]): Unit = {
     callbacks.foreach(_.startIteration(iter, weights))
   }
 
-  override def startMinibatch(iter : Int, weights : Weights[W], miniBatch : Array[MinibatchInput[T, W]]) : Unit = {
+  override def startMinibatch(iter: Int, weights: Weights[W],
+                              miniBatch: Array[MinibatchInput[T, W]]): Unit = {
     callbacks.foreach(_.startMinibatch(iter, weights, miniBatch))
   }
-  override def endMinibatch(iter : Int, weights : Weights[W], miniBatch : Array[MinibatchOutput[T, W]]) : Unit = {
+  override def endMinibatch(iter: Int, weights: Weights[W],
+                            miniBatch: Array[MinibatchOutput[T, W, S1, S2]]): Unit = {
     callbacks.foreach(_.endMinibatch(iter, weights, miniBatch))
   }
 
-  override def endIteration(iter : Int, weights : Weights[W]) : Unit = {
-    callbacks.foreach(_.endIteration(iter, weights))
+  override def endIteration(iter: Int, weights: Weights[W], state1: S1, state2: S2): Unit = {
+    callbacks.foreach(_.endIteration(iter, weights, state1, state2))
   }
 
-  override def objectiveValCheck(primal : Double, dual : Double, iter : Int, w : Weights[W]) : Unit = {
+  override def objectiveValCheck(primal: Double, dual: Double, iter: Int, w: Weights[W]): Unit = {
     callbacks.foreach(_.objectiveValCheck(primal, dual, iter, w))
   }
 
-  override def converged(weights : Weights[W], data : Seq[Instance[T, W]], iter : Int, numNewConstraints : Int) : Boolean = {
+  override def converged(weights: Weights[W], data: Seq[Instance[T, W]],
+                         iter: Int, numNewConstraints: Int): Boolean = {
     callbacks.forall(_.converged(weights, data, iter, numNewConstraints))
   }
 

src/main/scala/epic/logo/Decoder.scala

@@ -1,7 +1,8 @@
 package epic.logo
 
-trait Decoder[T, W] {
-
-  def decode(weights : Weights[W], instance : T) : (W, Double)
+trait Decoder[T, W, OracleS, MaxerS] {
 
+  def decode(weights : Weights[W], instance : T) : (W, Double, OracleS, MaxerS)
+  def initialState: (OracleS, MaxerS)
+  def reduceStates(states1: (OracleS, MaxerS), states2: (OracleS, MaxerS)): (OracleS, MaxerS)
 }

src/main/scala/epic/logo/ExpectationInferencer.scala

@@ -1,7 +1,6 @@
 package epic.logo
 
-trait ExpectationInferencer[T, W] {
-
-  def expectations(weights : Weights[W], instance : T) : (W, Double)
+trait ExpectationInferencer[T, W, S] extends Inferencer[S] {
 
+  def expectations(weights: Weights[W], instance: T): (W, Double, S)
 }

src/main/scala/epic/logo/Inferencer.scala

@@ -0,0 +1,7 @@
+package epic.logo
+
+trait Inferencer[S] {
+  def initialState: S
+  def reduceStates(state1: S, state2: S): S
+
+}

src/main/scala/epic/logo/IterationCallback.scala

@@ -1,17 +1,17 @@
 package epic.logo
 
-trait IterationCallback[T, W] {
+trait IterationCallback[T, W, OracleS, MaxerS] {
 
   def startIteration(iter : Int, weights : Weights[W]) : Unit = {}
 
   def startMinibatch(iter : Int, weights : Weights[W], miniBatch : Array[MinibatchInput[T, W]]) : Unit = {}
-  def endMinibatch(iter : Int, weights : Weights[W], miniBatch : Array[MinibatchOutput[T, W]]) : Unit = {}
-  def endIteration(iter : Int, weights : Weights[W]) : Unit = {}
+  def endMinibatch(iter : Int, weights : Weights[W], miniBatch : Array[MinibatchOutput[T, W, OracleS, MaxerS]]) : Unit = {}
+  def endIteration(iter : Int, weights : Weights[W], oracleState: OracleS, maxerState: MaxerS) : Unit = {}
 
   def objectiveValCheck(primal : Double, dual : Double, iter : Int, weights : Weights[W]) : Unit = {}
 
   def converged(weights : Weights[W], data : Seq[Instance[T, W]], iter : Int, numNewConstraints : Int) : Boolean = false
 
 }
 
-case class NullIterationCallback[T, W]() extends IterationCallback[T, W]
+case class NullIterationCallback[T, W, S1, S2]() extends IterationCallback[T, W, S1, S2]

src/main/scala/epic/logo/LogLikelihoodDecoder.scala

@@ -2,15 +2,19 @@ package epic.logo
 
 import breeze.math.MutableInnerProductModule
 
-class LogLikelihoodDecoder[T, W](
-    val inferencer: OracleInferencer[T, _, W], val summer: ExpectationInferencer[T, W])(
-        implicit space: MutableInnerProductModule[W, Double]) extends Decoder[T, W] {
+case class LogLikelihoodDecoder[T, W, OracleS, ExpectationS](
+    inferencer: OracleInferencer[T, W, OracleS], summer: ExpectationInferencer[T, W, ExpectationS])(
+        implicit space: MutableInnerProductModule[W, Double]) extends Decoder[T, W, OracleS, ExpectationS] {
   import space._
 
-  def decode(weights : Weights[W], instance : T) : (W, Double) = {
-    val (y_*, f_*, l_*) = inferencer.oracle(weights, instance)
-    val (f, l) = summer.expectations(weights, instance)
-    ((f_* - f), l - l_*)
+  def decode(weights : Weights[W], instance : T) : (W, Double, OracleS, ExpectationS) = {
+    val (y_*, f_*, l_*, state_*) = inferencer.oracle(weights, instance)
+    val (f, l, state) = summer.expectations(weights, instance)
+    ((f_* - f), l - l_*, state_*, state)
   }
 
+  def initialState = (inferencer.initialState, summer.initialState)
+  def reduceStates(states1: (OracleS, ExpectationS), states2: (OracleS, ExpectationS)): (OracleS, ExpectationS) =
+    (inferencer.reduceStates(states1._1, states2._1),
+      summer.reduceStates(states1._2, states2._2))
 }

src/main/scala/epic/logo/LossAugmentedArgmaxInferencer.scala

@@ -1,9 +1,9 @@
 package epic.logo
 
-trait LossAugmentedArgmaxInferencer[T, Y, W] extends ArgmaxInferencer[T, Y, W] {
+trait LossAugmentedArgmaxInferencer[T, W, S] extends ArgmaxInferencer[T, W, S] {
 
-  def argmax(weights : Weights[W], instance : T) : (Y, W, Double) = lossAugmentedArgmax(weights, instance, 1.0, 0.0)
+  def argmax(weights : Weights[W], instance : T) : (Y, W, Double, S) = lossAugmentedArgmax(weights, instance, 1.0, 0.0)
 
-  def lossAugmentedArgmax(weights : Weights[W], instance : T, weightsWeight : Double, lossWeight : Double) : (Y, W, Double)
+  def lossAugmentedArgmax(weights : Weights[W], instance : T, weightsWeight : Double, lossWeight : Double) : (Y, W, Double, S)
 
 }

src/main/scala/epic/logo/LossAugmentedExpectationInferencer.scala

@@ -1,9 +1,9 @@
 package epic.logo
 
-trait LossAugmentedExpectationInferencer[T, W] extends ExpectationInferencer[T, W] {
+trait LossAugmentedExpectationInferencer[T, W,  S] extends ExpectationInferencer[T, W, S] {
 
-  def expectations(weights : Weights[W], instance : T) : (W, Double) = lossAugmentedExpectations(weights, instance, 1.0, 0.0)
+  def expectations(weights : Weights[W], instance : T) : (W, Double, S) = lossAugmentedExpectations(weights, instance, 1.0, 0.0)
 
-  def lossAugmentedExpectations(weights : Weights[W], instance : T, weightsWeight : Double, lossWeight : Double) : (W, Double)
+  def lossAugmentedExpectations(weights : Weights[W], instance : T, weightsWeight : Double, lossWeight : Double) : (W, Double, S)
 
 }

src/main/scala/epic/logo/LossAugmentedMaxMarginDecoder.scala

@@ -2,13 +2,21 @@ package epic.logo
 
 import breeze.math.MutableInnerProductModule
 
-class LossAugmentedMaxMarginDecoder[T, W](val oracleInferencer: OracleInferencer[T, _, W], val argmaxer: LossAugmentedArgmaxInferencer[T, _, W])(implicit space: MutableInnerProductModule[W, Double]) extends Decoder[T, W] {
+case class LossAugmentedMaxMarginDecoder[T, W, OracleS, ArgmaxerS](
+    oracleInferencer: OracleInferencer[T, W, OracleS],
+    argmaxer: LossAugmentedArgmaxInferencer[T, W, ArgmaxerS])(implicit space: MutableInnerProductModule[W, Double])
+    extends Decoder[T, W, OracleS, ArgmaxerS] {
   import space._
 
-  def decode(weights: Weights[W], instance: T): (W, Double) = {
-    val (y_*, f_*, l_*) = oracleInferencer.oracle(weights, instance)
-    val (y, f, l) = argmaxer.lossAugmentedArgmax(weights, instance, 1.0, 1.0)
-    (f_* - f, l - l_*)
+  def decode(weights: Weights[W], instance: T): (W, Double, OracleS, ArgmaxerS) = {
+    val (y_*, f_*, l_*, state_*) = oracleInferencer.oracle(weights, instance)
+    val (y, f, l, state) = argmaxer.lossAugmentedArgmax(weights, instance, 1.0, 1.0)
+    (f_* - f, l - l_*, state_*, state)
   }
 
+  def initialState: (OracleS, ArgmaxerS) = (oracleInferencer.initialState, argmaxer.initialState)
+  def reduceStates(states1: (OracleS, ArgmaxerS), states2: (OracleS, ArgmaxerS)): (OracleS, ArgmaxerS) =
+    (oracleInferencer.reduceStates(states1._1, states2._1),
+      argmaxer.reduceStates(states1._2, states2._2))
+
 }

src/main/scala/epic/logo/MaxMarginDecoder.scala

@@ -2,13 +2,21 @@ package epic.logo
 
 import breeze.math.MutableInnerProductModule
 
-class MaxMarginDecoder[T, W](val oracleInferencer : OracleInferencer[T, _, W], val argmaxer : ArgmaxInferencer[T, _, W])(implicit space: MutableInnerProductModule[W, Double]) extends Decoder[T, W] {
+case class MaxMarginDecoder[T, W, OracleS, MaxerS](
+  oracleInferencer: OracleInferencer[T, W, OracleS],
+  argmaxer: ArgmaxInferencer[T, W, MaxerS])(implicit space: MutableInnerProductModule[W, Double])
+    extends Decoder[T, W, OracleS, MaxerS] {
   import space._
 
-  def decode(weights : Weights[W], instance : T) : (W, Double) = {
-    val (y_*, f_*, l_*) = oracleInferencer.oracle(weights, instance)
-    val (y, f, l) = argmaxer.argmax(weights, instance)
-    (f_* - f, l - l_*)
+  def decode(weights : Weights[W], instance : T) : (W, Double, OracleS, MaxerS) = {
+    val (y_*, f_*, l_*, state_*) = oracleInferencer.oracle(weights, instance)
+    val (y, f, l, state) = argmaxer.argmax(weights, instance)
+    (f_* - f, l - l_*, state_*, state)
   }
 
+  def initialState: (OracleS, MaxerS) = (oracleInferencer.initialState, argmaxer.initialState)
+  def reduceStates(states1: (OracleS, MaxerS), states2: (OracleS, MaxerS)): (OracleS, MaxerS) =
+    (oracleInferencer.reduceStates(states1._1, states2._1),
+      argmaxer.reduceStates(states1._2, states2._2))
+
 }

src/main/scala/epic/logo/MaxMarginRankingDecoder.scala

@@ -2,12 +2,19 @@ package epic.logo
 
 import breeze.math.MutableInnerProductModule
 
-class MaxMarginRankingDecoder[T, W](val inferencer : LossAugmentedArgmaxInferencer[T, _, W], val gamma : Double = 0.0)(implicit space: MutableInnerProductModule[W, Double]) extends Decoder[T, W] {
-import space._
-  def decode(weights: Weights[W], instance: T): (W, Double) = {
-    val (y_min, df_min, l_min) = inferencer.lossAugmentedArgmax(weights, instance, -1.0, -(1.0 + gamma))
-    val (y_max, df_max, l_max) = inferencer.lossAugmentedArgmax(weights, instance, 1.0, 1.0)
-    (df_min - df_max, l_max - (1.0 + gamma) * l_min)
+case class MaxMarginRankingDecoder[T, W, S](
+    inferencer: LossAugmentedArgmaxInferencer[T, W, S],
+    gamma: Double = 0.0)(implicit space: MutableInnerProductModule[W, Double]) extends Decoder[T, W, S, S] {
+  import space._
+  def decode(weights: Weights[W], instance: T): (W, Double, S, S) = {
+    val (y_min, df_min, l_min, s_min) = inferencer.lossAugmentedArgmax(weights, instance, -1.0, -(1.0 + gamma))
+    val (y_max, df_max, l_max, s_max) = inferencer.lossAugmentedArgmax(weights, instance, 1.0, 1.0)
+    (df_min - df_max, l_max - (1.0 + gamma) * l_min, s_min, s_max)
   }
 
+  def initialState: (S, S) = (inferencer.initialState, inferencer.initialState)
+  def reduceStates(states1: (S, S), states2: (S, S)): (S, S) =
+    (inferencer.reduceStates(states1._1, states2._1),
+      inferencer.reduceStates(states1._2, states2._2))
+
 }

src/main/scala/epic/logo/MinibatchInput.scala

@@ -1,5 +1,7 @@
 package epic.logo
 
-case class MinibatchInput[T, W](val instance : Instance[T, W], val instanceNum : Int)
+case class MinibatchInput[T, W](instance : Instance[T, W], instanceNum : Int)
 
-case class MinibatchOutput[T, W](val instance : Instance[T, W], val instanceNum : Int, val df : W, val loss : Double)
+case class MinibatchOutput[T, W, OracleS, MaxerS](
+  instance: Instance[T, W], instanceNum: Int, df: W, loss: Double,
+  oracleState: OracleS, maxerState: MaxerS)

src/main/scala/epic/logo/MulticlassLossAugmentedArgmaxInferencer.scala

@@ -3,16 +3,18 @@ import breeze.util.Index
 import breeze.math.MutableInnerProductModule
 
 class MulticlassLossAugmentedArgmaxInferencer[L, F, W](validLabels: IndexedSeq[L], labelConjoiner: (L, F) => W)
-(implicit space: MutableInnerProductModule[W, Double])
-    extends LossAugmentedArgmaxInferencer[LabeledDatum[L, F], L, W] {
+    extends LossAugmentedArgmaxInferencer[LabeledDatum[L, F], W, Unit] {
+  type Y = L
 
   def lossAugmentedArgmax(weights: Weights[W], instance: LabeledDatum[L, F],
-                          weightsWeight: Double, lossWeight: Double): (L, W, Double) = {
+                          weightsWeight: Double, lossWeight: Double): (L, W, Double, Unit) = {
     validLabels.map(label => {
       val loss = if (instance.label == null || label.equals(instance.label)) 0.0 else 1.0
       val labeledFeatureVector = labelConjoiner(label, instance.features)
-      (label,labeledFeatureVector, loss)
-    }).maxBy { case (label, features, loss) => weightsWeight * (weights * features) + lossWeight * loss }
+      (label,labeledFeatureVector, loss, ())
+    }).maxBy { case (label, features, loss, _) => weightsWeight * (weights * features) + lossWeight * loss }
   }
 
+  def initialState: Unit = ()
+  def reduceStates(a: Unit, b: Unit): Unit = ()
 }

src/main/scala/epic/logo/MulticlassOneSlackLossAugmentedArgmaxInferencer.scala

@@ -3,16 +3,24 @@ package epic.logo
 import breeze.math.MutableInnerProductModule
 
 class MulticlassOneSlackLossAugmentedArgmaxInferencer[L, F, W](
-    inferencer: MulticlassLossAugmentedArgmaxInferencer[L, F, W],
-    emptyFeatureVector: W)(implicit space: MutableInnerProductModule[W, Double]) extends LossAugmentedArgmaxInferencer[Seq[LabeledDatum[L, F]], Seq[L], W] {
+  inferencer: MulticlassLossAugmentedArgmaxInferencer[L, F, W],
+  emptyFeatureVector: W)(implicit space: MutableInnerProductModule[W, Double])
+    extends LossAugmentedArgmaxInferencer[Seq[LabeledDatum[L, F]], W, Unit] {
+  override type Y = Seq[L]
   import space._
 
   def lossAugmentedArgmax(weights: Weights[W], instance: Seq[LabeledDatum[L, F]], weightsWeight: Double,
-                          lossWeight: Double): (Seq[L], W, Double) = {
-    instance.map(i => inferencer.lossAugmentedArgmax(weights, i, weightsWeight, lossWeight))
+                          lossWeight: Double): (Seq[L], W, Double, Unit) = {
+   val (labels, fvs, losses) =
+     instance.map(i => inferencer.lossAugmentedArgmax(weights, i, weightsWeight, lossWeight))
       .foldLeft((Seq.empty[L], emptyFeatureVector, 0.0)) { (acc, curr) =>
         (acc._1 :+ curr._1, acc._2 + curr._2, acc._3 + curr._3)
       }
+   (labels, fvs, losses, ())
   }
 
+  def initialState: Unit = ()
+
+  def reduceStates(s1: Unit, s2: Unit): Unit = ()
+
 }

src/main/scala/epic/logo/MulticlassOneSlackOracleInferencer.scala

@@ -4,15 +4,21 @@ import breeze.math.MutableInnerProductModule
 
 class MulticlassOneSlackOracleInferencer[L, F, W](
   val inferencer: MulticlassOracleInferencer[L, F, W])(implicit space: MutableInnerProductModule[W, Double])
-    extends OracleInferencer[Seq[LabeledDatum[L, F]], Seq[L], W] {
+    extends OracleInferencer[Seq[LabeledDatum[L, F]], W, Unit] {
   import space._
+  type Y = Seq[L]
 
-  def oracle(weights: Weights[W], instance: Seq[LabeledDatum[L, F]]): (Seq[L], W, Double) = {
+  def oracle(weights: Weights[W], instance: Seq[LabeledDatum[L, F]]): (Seq[L], W, Double, Unit) = {
     val zeroVector = space.zeroLike(inferencer.labelConjoiner(instance.head.label, instance.head.features))
-    instance.map(i => inferencer.oracle(weights, i))
+    val (labels, fvs, losses) = instance.map(i => inferencer.oracle(weights, i))
       .foldLeft((Seq.empty[L], zeroVector, 0.0)) { (acc, curr) =>
         (acc._1 :+ curr._1, acc._2 + curr._2, acc._3 + curr._3)
       }
+    (labels, fvs, losses, ())
   }
 
+  def initialState: Unit = ()
+
+  def reduceStates(state1: Unit, state2: Unit): Unit = ()
+
 }

src/main/scala/epic/logo/MulticlassOracleInferencer.scala

@@ -5,11 +5,15 @@ import breeze.math.MutableInnerProductModule
 
 case class MulticlassOracleInferencer[L, F, W](
   validLabels: IndexedSeq[L], labelConjoiner: (L, F) => W)(implicit space: MutableInnerProductModule[W, Double])
-    extends OracleInferencer[LabeledDatum[L, F], L, W] {
+    extends OracleInferencer[LabeledDatum[L, F], W, Unit] {
+  override type Y = L
 
-  def oracle(weights : Weights[W], instance : LabeledDatum[L, F]) : (L, W, Double) = {
-    (instance.label, labelConjoiner(instance.label, instance.features), 0.0)
+  def oracle(weights : Weights[W], instance : LabeledDatum[L, F]) : (L, W, Double, Unit) = {
+    (instance.label, labelConjoiner(instance.label, instance.features), 0.0, ())
 
   }
 
+  def initialState = ()
+  def reduceStates(a: Unit, b: Unit) = ()
+
 }

src/main/scala/epic/logo/ObjectiveFunctionConvergenceChecker.scala

@@ -4,7 +4,7 @@ import scala.collection.Seq
 
 case class ObjectiveFunctionConvergenceChecker[W](
     objective: ObjectiveFunction[W], maxNumIters: Int,
-    callback: IterationCallback[_, W], tol: Double) extends ConvergenceChecker[W] {
+    callback: IterationCallback[_, W, _, _], tol: Double) extends ConvergenceChecker[W] {
 
   def converged(weights: Weights[W], data: Seq[Instance[_, W]], iter: Int, numNewConstraints: Int): Boolean = {
     iter >= maxNumIters || (objectiveConverged(weights, data, iter) && numNewConstraints == 0)
@@ -13,7 +13,7 @@ case class ObjectiveFunctionConvergenceChecker[W](
   private def objectiveConverged(weights: Weights[W], data: Seq[Instance[_, W]], iter: Int) = {
     val (primal, dual) = objective.calculatePrimalAndDual(weights, data)
     callback.objectiveValCheck(primal, dual, iter, weights)
-    NumUtils.approxEquals(primal, dual, 1e-10)
+    NumUtils.approxEquals(primal, dual, tol)
   }
 
 }

src/main/scala/epic/logo/OracleInferencer.scala

@@ -1,7 +1,7 @@
 package epic.logo
 
-trait OracleInferencer[T, Y, W] {
-
-  def oracle(weights : Weights[W], instance : T) : (Y, W, Double)
+trait OracleInferencer[T, W, S] extends Inferencer[S] {
+  type Y
 
+  def oracle(weights : Weights[W], instance : T) : (Y, W, Double, S)
 }