diff --git a/compiler/modules/write_mask_and_array.py b/compiler/modules/write_mask_and_array.py
index 75d39a94..c54d66c2 100644
--- a/compiler/modules/write_mask_and_array.py
+++ b/compiler/modules/write_mask_and_array.py
@@ -63,8 +63,9 @@ class write_mask_and_array(design.design):
 
     def add_modules(self):
         # Size the AND gate for the number of write drivers it drives, which is equal to the write size.
+        # Assume stage effort of 3 to compute the size
         self.and2 = factory.create(module_type="pand2",
-                                   size=self.write_size)
+                                   size=self.write_size/4.0)
         self.add_mod(self.and2)
 
 
@@ -92,18 +93,12 @@ class write_mask_and_array(design.design):
         else:
             self.driver_spacing = self.driver.width
 
-        if (self.words_per_row == 1):
-            wmask_en_len = (self.write_size * self.driver_spacing)
-            if self.driver_spacing * self.write_size < self.and2.width:
-                debug.error("Cannot layout write mask AND array. One pand2 is longer than the corresponding write drivers.")
-        else:
-            wmask_en_len = self.words_per_row * (self.write_size * self.driver_spacing)
-            if wmask_en_len < self.and2.width:
-                debug.error("Cannot layout write mask AND array. One pand2 is longer than the corresponding write drivers.")
-        self.wmask_en_len = wmask_en_len
+        self.wmask_en_len = self.words_per_row * (self.write_size * self.driver_spacing)
+        debug.check(self.wmask_en_len >= self.and2.width,
+                    "Write mask AND is wider than the corresponding write drivers {0} vs {1}.".format(self.and2.width,self.wmask_en_len))
 
         for i in range(self.num_wmasks):
-            base = vector(i * wmask_en_len, 0)
+            base = vector(i * self.wmask_en_len, 0)
             self.and2_insts[i].place(base)
 
 
diff --git a/compiler/pgates/pand2.py b/compiler/pgates/pand2.py
index b241dc55..d71b1e92 100644
--- a/compiler/pgates/pand2.py
+++ b/compiler/pgates/pand2.py
@@ -36,7 +36,6 @@ class pand2(pgate.pgate):
         self.nand = factory.create(module_type="pnand2",height=self.height) 
         self.add_mod(self.nand)
 
-        # Assume stage effort of 3
         self.inv = factory.create(module_type="pdriver", neg_polarity=True, fanout=3*self.size, height=self.height)
         self.add_mod(self.inv)
 
diff --git a/compiler/pgates/pdriver.py b/compiler/pgates/pdriver.py
index 36811a5e..ec55f0c7 100644
--- a/compiler/pgates/pdriver.py
+++ b/compiler/pgates/pdriver.py
@@ -45,7 +45,7 @@ class pdriver(pgate.pgate):
             self.num_stages = len(self.size_list)
         else:
             # Find the optimal number of stages for the given effort
-            self.num_stages = max(1,int(round(log(self.fanout)/log(self.stage_effort))))
+            self.num_stages = max(1,int(round(self.fanout**(1/self.stage_effort))))
 
             # Increase the number of stages if we need to fix polarity
             if self.neg_polarity and (self.num_stages%2==0):