--- a/pylearn/algorithms/mcRBM.py	Tue Aug 24 14:54:41 2010 -0400
+++ b/pylearn/algorithms/mcRBM.py	Tue Aug 24 15:28:13 2010 -0400
@@ -309,6 +309,15 @@
 
 class MeanCovRBM(object):
     """Container for mcRBM parameters that gives more convenient access to mcRBM methods.
+
+    Attributes:
+
+      - U - the covariance filters
+      - W - the mean filters
+      - a - the visible bias
+      - b - the covariance bias
+      - c - the mean bias
+    
     """
 
     params = property(lambda s: [s.U, s.W, s.a, s.b, s.c])
@@ -316,40 +325,36 @@
     n_visible = property(lambda s: s.W.value.shape[0])
 
     def __init__(self, U, W, a, b, c):
-        self.U = as_shared(U, 'U')
-        self.W = as_shared(W, 'W')
-        self.a = as_shared(a, 'a')
-        self.b = as_shared(b, 'b')
-        self.c = as_shared(c, 'c')
-
-        assert self.b.type.dtype == 'float32'
+        self.__dict__.update(locals())
+        del self.__dict__['self']
 
     @classmethod
-    def new_from_dims(cls, 
-            n_I,  # input dimensionality
-            n_K,  # number of covariance hidden units
-            n_F,  # number of covariance filters (squared)
-            n_J,  # number of mean filters (linear)
-            seed = 8923402190,
-            ):
+    def new_from_dims(cls, n_I, n_K, n_J, rng = 8923402190):
         """
         Return a MeanCovRBM instance with randomly-initialized parameters.
-        """
-        rng = np.random.RandomState(seed)
 
-        # initialization taken from Marc'Aurelio
+        :param n_I: input dimensionality
+        :param n_K: number of covariance hidden units
+        :param n_J: number of mean filters (linear)
+        :param rng: seed or numpy RandomState object to initialize params
+        """
+        if not hasattr(rng, 'randn'):
+            rng = np.random.RandomState(rng)
 
+        def shrd(X,name):
+            return shared(X.astype(floatX), name=name)
+
+        # initialization taken from train_mcRBM.py
         return cls(
-                U = 0.02 * rng.randn(n_I, n_F),
-                W = 0.05 * rng.randn(n_I, n_J),
-                #W = rng.randn(n_I, n_J)/np.sqrt((n_I+n_J)/2),
-                a = np.ones(n_I)*(0),
-                b = np.ones(n_K)*2,
-                c = np.ones(n_J)*(-2),)
+                U = shrd(0.02 * rng.randn(n_I, n_K),'U'),
+                W = shrd(0.05 * rng.randn(n_I, n_J),'W'),
+                a = shrd(np.ones(n_I)*(0),'a'),
+                b = shrd(np.ones(n_K)*2,'b'),
+                c = shrd(np.ones(n_J)*(-2),'c'))
 
     def __getstate__(self):
-        # unpack shared containers, which may have references to Theano stuff
-        # and are not a long-term stable data type.
+        # unpack shared containers, which may have references to Theano stuff and are not a
+        # long-term stable data type.
         return dict(
                 U = self.U.value,
                 W = self.W.value,
@@ -357,7 +362,10 @@
                 c = self.c.value)
 
     def __setstate__(self, dct):
-        self.__init__(**dct) # calls as_shared on pickled arrays
+        d = dict(dct)
+        for key in ['U', 'W', 'a', 'b', 'c']:
+            d[key] = shared(d[key], name=key)
+        self.__init__(**d)
 
     def hmc_sampler(self, n_particles, rng=7823748):
         """Return an HMC_sampler that will draw samples from the distribution over visible
@@ -419,11 +427,7 @@
     n_F=256
     n_J=100
 
-    rbm = MeanCovRBM.new_from_dims(n_I=R*C,
-            n_K=n_K,
-            n_J=n_J, 
-            n_F=n_F,
-            ) 
+    rbm = MeanCovRBM.new_from_dims(n_I=R*C, n_K=n_K, n_J=n_J) 
 
     sampler = rbm.hmc_sampler(n_particles=batchsize)