diff --git a/lagrangian.cabal b/lagrangian.cabal
--- a/lagrangian.cabal
+++ b/lagrangian.cabal
@@ -10,7 +10,7 @@
 -- PVP summary:      +-+------- breaking API changes
 --                   | | +----- non-breaking API additions
 --                   | | | +--- code changes with no API change
-version:             0.4.0.1
+version:             0.5.0.0
 
 -- A short (one-line) description of the package.
 synopsis:            Solve Lagrange multiplier problems
@@ -94,11 +94,11 @@
   build-depends: base ==4.6.*,
                  nonlinear-optimization ==0.3.*, 
                  vector ==0.10.*, 
-                 ad ==3.3.*,
+                 ad ==3.4.*,
                  hmatrix == 0.14.*, 
-                 test-framework ==0.6.*, 
-                 test-framework-hunit ==0.2.*, 
-                 test-framework-quickcheck2 ==0.2.*,
+                 test-framework ==0.8.*, 
+                 test-framework-hunit ==0.3.*, 
+                 test-framework-quickcheck2 ==0.3.*,
                  HUnit == 1.2.*
 
 
diff --git a/src/Numeric/AD/Lagrangian.hs b/src/Numeric/AD/Lagrangian.hs
--- a/src/Numeric/AD/Lagrangian.hs
+++ b/src/Numeric/AD/Lagrangian.hs
@@ -11,6 +11,7 @@
 module Numeric.AD.Lagrangian (
     -- *** Helper types
     AD2,
+    FU(..),
     (<=>),
     Constraint,
     -- ** Solver
@@ -18,4 +19,5 @@
     minimize,
     -- *** Experimental features
     feasible) where
-import Numeric.AD.Lagrangian.Internal (AD2, (<=>), maximize, minimize, feasible, Constraint)
+import Numeric.AD.Lagrangian.Internal (AD2, FU(..), 
+    (<=>), maximize, minimize, feasible, Constraint)
diff --git a/src/Numeric/AD/Lagrangian/Internal.hs b/src/Numeric/AD/Lagrangian/Internal.hs
--- a/src/Numeric/AD/Lagrangian/Internal.hs
+++ b/src/Numeric/AD/Lagrangian/Internal.hs
@@ -13,24 +13,25 @@
 import Numeric.AD.Internal.Tower
 
 infixr 1 <=>
--- | This is just a little bit of sugar for (,) to make constraints look like 
---  equals
-(<=>) :: ([a] -> a) -> a -> Constraint a
-g <=> c = (g,c)
--- | The type for the contraints.
---   Given a constraint @g(x, y, ...) = c@, we would represent it as @(g, c)@.
---   or with sugar @g@ '<=>' @c@
-type Constraint a = ([a] -> a, a)
+-- | Build a 'Constraint' from a function and a constant
+(<=>) :: (forall s r. (Mode s, Mode r) => [AD2 s r a] -> AD2 s r a) -> a -> Constraint a
+g <=> c = (FU g,c)
 
-type AD2 s r a = AD s (AD r Double)
+-- | A constraint of the form @g(x, y, ...) = c@. See '<=>' for constructing a 'Constraint'.
+type Constraint a = (FU a, a)
 
+type AD2 s r a = AD s (AD r a)
+
+-- | A newtype wrapper for working with the rank 2 types constraint functions. 
+newtype FU a = FU {unFU :: forall s r. (Mode s, Mode r) => [AD2 s r a] -> AD2 s r a}
+
 -- | This is the lagrangian multiplier solver. It is assumed that the 
 --   objective function and all of the constraints take in the 
 --   same amount of arguments.
 minimize :: Double
       -> (forall s r. (Mode s, Mode r) => [AD2 s r Double] -> AD2 s r Double) 
         -- ^ The function to minimize
-      -> (forall s r. (Mode s, Mode r) => [Constraint (AD2 s r Double)] ) 
+      -> [Constraint Double]
       -- ^ The constraints as pairs @g \<=\> c@ which represent equations 
       --   of the form @g(x, y, ...) = c@
       -> Int 
@@ -44,8 +45,7 @@
     obj argsAndLams = 
         squaredGrad (lagrangian toMin constraints argCount) argsAndLams
     
-    -- The mode does matter but I need to add annotation for the type checker
-    constraintCount = length (constraints :: [Constraint (AD Tower (AD Tower Double))])
+    constraintCount = length constraints 
     
     -- perhaps this should be exposed
     guess = U.replicate (argCount + constraintCount) (1.0 :: Double) 
@@ -66,7 +66,7 @@
 maximize :: Double
       -> (forall s r. (Mode s, Mode r) => [AD2 s r Double] -> AD2 s r Double) 
         -- ^ The function to maximize
-      -> (forall s r. (Mode s, Mode r) => [Constraint (AD2 s r Double)] ) 
+      -> [Constraint Double] 
       -- ^ The constraints as pairs @g \<=\> c@ which represent equations 
       --   of the form @g(x, y, ...) = c@
       -> Int 
@@ -78,18 +78,18 @@
 maximize tolerance toMax constraints argCount = 
     minimize tolerance (negate1 . toMax) constraints argCount
 
-lagrangian :: Num a 
-           => ([a] -> a)
+lagrangian :: (Num a, Mode s, Mode r)
+           => (forall s r. (Mode s, Mode r) => [AD2 s r a] -> AD2 s r a) 
            -> [Constraint a]
            -> Int
-           -> [a]  
-           -> a
+           -> [AD2 s r a]  
+           -> AD2 s r a
 lagrangian f constraints argCount argsAndLams = result where
     args = take argCount argsAndLams
     lams = drop argCount argsAndLams
     
     -- g(x, y, ...) = c <=> g(x, y, ...) - c = 0
-    appliedConstraints = fmap (\(f, c) -> f args - c) constraints
+    appliedConstraints = fmap (\(FU f, c) -> f args - (auto . auto) c) constraints
     
     -- L(x, y, ..., lam0, ...) = f(x, y, ...) + lam0 * (g0 - c0) ... 
     result = f args + (sum . zipWith (*) lams $ appliedConstraints)
@@ -104,7 +104,7 @@
 --   exactly how to implement that. This just checks the feasiblility at a point.
 --   If this ever returns false, 'solve' can fail.
 feasible :: (forall s r. (Mode s, Mode r) => [AD2 s r Double] -> AD2 s r Double)
-         -> (forall s r. (Mode s, Mode r) => [Constraint (AD2 s r Double)] )
+         -> [Constraint Double]
          -> [Double]
          -> Bool
 feasible toMin constraints points = result where
diff --git a/tests/Main.hs b/tests/Main.hs
--- a/tests/Main.hs
+++ b/tests/Main.hs
@@ -23,8 +23,10 @@
 --class Approximate a where
 --    x =~= y :: a -> a -> Bool
 
+
+
 entropyTest = (S.sum . S.map abs $ S.zipWith (-) actual expected) < 0.02 @?= True  where
-    Right actual = fst <$> maximize 0.00001 f [(\xs -> sum xs, 1)] 3
+    Right actual = fst <$> maximize 0.00001 f [sum <=> 1] 3
     expected  = S.fromList [0.33, 0.33, 0.33]
     f :: Floating a => [a] -> a
     f = negate . sum . map (\x -> x * log x)