--- a/pylearn/gd/__init__.py	Wed Apr 06 13:52:34 2011 -0400
+++ b/pylearn/gd/__init__.py	Wed Apr 06 13:53:00 2011 -0400
@@ -9,3 +9,5 @@
  - Stopping criteria (incl. for use in theano functions)
 
 """
+
+from sgd import sgd_updates, sgd_momentum_updates

--- a/pylearn/gd/sgd.py	Wed Apr 06 13:52:34 2011 -0400
+++ b/pylearn/gd/sgd.py	Wed Apr 06 13:53:00 2011 -0400
@@ -18,9 +18,30 @@
         iter(stepsizes)
     except:
         stepsizes = [stepsizes for p in params]
+    if len(params) != len(grads):
+        raise ValueError('params and grads have different lens')
     updates = [(p, p - step * gp) for (step, p, gp) in zip(stepsizes, params, grads)]
     return updates
 
+def sgd_momentum_updates(params, grads, stepsizes, momentum=0.9):
+    # if stepsizes is just a scalar, expand it to match params
+    try:
+        iter(stepsizes)
+    except:
+        stepsizes = [stepsizes for p in params]
+    try:
+        iter(momentum)
+    except:
+        momentum = [momentum for p in params]
+    if len(params) != len(grads):
+        raise ValueError('params and grads have different lens')
+    headings = [theano.shared(p.get_value(borrow=False)*0) for p in params]
+    updates = []
+    for s, p, gp, m, h in zip(stepsizes, params, grads, momentum, headings):
+        updates.append((p, p + s * h))
+        updates.append((h, m*h - (1-m)*gp))
+    return updates
+
 
 class StochasticGradientDescent(theano.Module):
     """Fixed stepsize gradient descent