--- a/pylearn/algorithms/mcRBM.py	Tue Aug 24 14:03:07 2010 -0400
+++ b/pylearn/algorithms/mcRBM.py	Tue Aug 24 14:08:24 2010 -0400
@@ -269,13 +269,30 @@
 
     """
     U, W, a, b, c = rbm
-
     t0 = -TT.sum(TT.nnet.softplus(hidden_cov_units_preactivation_given_v(rbm, v)),axis=1)
     t1 = -TT.sum(TT.nnet.softplus(c + dot(v,W)), axis=1)
     t2 =  0.5 * TT.sum(v**2, axis=1)
     t3 = -TT.dot(v, a)
     return t0 + t1 + t2 + t3, (t0, t1, t2, t3)
 
+def contrastive_gradient(rbm, pos_v, neg_v, U_l1_penalty=0, W_l1_penalty=0):
+    """Return a list of gradient expressions for the rbm parameters
+
+    :param pos_v: positive-phase sample of visible units
+    :param neg_v: negative-phase sample of visible units
+    :param U_l1_penalty: a scalar-valued multiplier on the L1 penalty on U
+    :param W_l1_penalty: a scalar-valued multiplier on the L1 penalty on W
+    """
+    U, W, a, b, c = rbm
+    pos_FE = free_energy_given_v(rbm, pos_v)
+    neg_FE = free_energy_given_v(rbm, neg_v)
+    c0 = (pos_FE - neg_FE).sum()
+    c1 = abs(U).sum()*U_l1_penalty 
+    c2 = abs(W).sum()*W_l1_penalty 
+    cost = c0 + c1 + c2
+    rval = theano.tensor.grad(cost, list(rbm))
+    return rval
+
 def expected_h_g_given_v(P, U, W, b, c, v):
     """Returns theano expression conditional expectations (`h`, `g`) in an mcRBM.
     
@@ -365,21 +382,12 @@
         else:
             return rval[0]
 
-    def contrastive_gradient(self, pos_v, neg_v, U_l1_penalty=0, W_l1_penalty=0):
+    def contrastive_gradient(self, *args, **kwargs)
         """Return a list of gradient expressions for self.params
 
-        :param pos_v: positive-phase sample of visible units
-        :param neg_v: negative-phase sample of visible units
+        See `contrastive_gradient` for parameters.
         """
-        pos_FE = self.free_energy_given_v(pos_v)
-        neg_FE = self.free_energy_given_v(neg_v)
-
-        gpos_FE = theano.tensor.grad(pos_FE.sum(), self.params)
-        gneg_FE = theano.tensor.grad(neg_FE.sum(), self.params)
-        rval = [ gp - gn for (gp,gn) in zip(gpos_FE, gneg_FE)]
-        rval[0] = rval[0] - TT.sign(self.U)*U_l1_penalty
-        rval[1] = rval[1] - TT.sign(self.W)*W_l1_penalty
-        return rval
+        return contrastive_gradient(self.params, *args, **kwargs)
 
 
 if __name__ == '__main__':