Function Composition

[pharringtonp19]

In Haskell, we can define function composition using . as in the following

f x = x^2 
g x = sqrt x 
m = f . g 

Is it possible to do the same thing in Dex?

def f (x:Float32) : Float32 = 
    x*x

def g (x:Float32) : Float32 = 
    sqrt x

m = g . f 

Error Message

Type error:
Expected: (a => b)
  Actual: (Float32 -> Float32)
(Solving for: [a, b])

m = g . f 
    ^^

[duvenaud]

Yes, it's straightforward to define a compose function:

def compose {a b c} (f : a -> b) (g : b -> c) : a -> c =
  \x. g (f x)

(compose sin cos) 1.0
0.666367

You can't use the f . g syntax only because . is already defined as indexing an array. But if you can find another symbol you like, you can define it as an infix operator like so:

def (...) {a b c} (f : a -> b) (g : b -> c) : a -> c =
  \x. g (f x)

(sin ... cos) 1.0
0.666367

[pharringtonp19]

@duvenaud That makes sense -- Thanks!

I am surprised that the following doesn't work, though.

def compose {a b c} (f : a -> b) (g : b -> c) : a -> c =
  \x. g (f x)

m = compose sin cos 

Error

Type error:Ambiguous type variable

m = compose sin cos
    ^^^^^^^^

I would assume, though, that I should be able to reason through this error once I understand the signature of cos

:t cos
((a:Type) ?-> (Floating a) ?=> a -> a)

[duvenaud]

Oh, I think that's a bug / limitation of the current type system. At least, I think it's the same as issue #867. I think the problem is that there are multiple possible floating point types, and when you write m = compose sin cos, the compiler isn't sure which one you want, and can't automatically build an m that's polymorphic in a.

You can do what you want by telling Dex which kind of number you want:

m : Float32 -> Float32 = compose sin cos
m 1.0
0.666367

This is annoying, although it usually isn't an issue when you're not at the top level, so it rarely comes up in practice for me.

[pharringtonp19]

@duvenaud That makes sense. Thanks for the suggestion/correction

[apaszke]

Yes, that's a fine point that differentiates Haskell's type system from Dex's. Haskell does let-generalization, meaning that if the type variables end up unconstrained, it will generalize them. Dex doesn't and complains about those variables being ambiguous. The reason why we err on the side of being conservative is because of programs like this one:

z = (polymorphic_expensive_array_factory : (Ix n) ?=> (n=>Int))
a = (for i:(Fin 1000). z.i + 2)
b = (for i:(Fin 1000). z.i + 2)

Of course it's a silly example, but if we did let-generalization in this case the z = zero assignment would do no work and polymorphic_expensive_array_factory would get inlined into body of a and b separately! But this is a performance footgun just waiting to fire!

[pharringtonp19]

@duvenaud

Having defined composition like you did above (I've used . to denote composition below), is it possible to introduce an additional operator that can define repeated composition? As in something like the following?

3 (.) f = f . f . f

How might we define (.) in dex?

[duvenaud]

Like so:

def compose_n_times {a} (n:Nat) (f:a->a) : a->a =
  \x. yield_state x \ansRef.
    for i:(Fin n).
      ansRef := f (get ansRef)

(compose_n_times 7 (\x. x + 1)) 100
107

This is pretty similar to what fold and scan do.

[pharringtonp19]

@duvenaud Thanks for the explanation. Much appreciated.

[pharringtonp19]

It's really nice that with dex we can write a training loop by composing a step function as follows:

(1000 *.* step) init_state

'### Function Composition

def (.*.) {a:Type} (f : a -> a) (g : a -> a)  :  (a -> a) =
    \x. g (f x)

def (*.*) {a:Type} (n:Nat)  (f : a -> a) : a -> a = 
    \x. yield_state x \ansRef. 
        for i:(Fin n).
            ansRef := f (get ansRef)
'### Objective Function 

def objective (x: Float) : Float = 
    log (x*x + 1.0 + sin(3*x)) + 1.5
'### Update/Step Function

def update (f: Float->Float) (lr:Float) (mom:Float) (state: (Float & Float)):  (Float & Float) =  
    v = fst state
    x = snd state 
    v = mom*v + grad f x
    (v,  x - lr * v)

'### Training

-- Hyperparameters
lr = 0.1 
mom = 0.9
init_state = (0.0, (-4.0))

-- Step Function
step = update objective lr mom  

-- Training "Loop"
opt_param = snd $ (1000 *.* step) init_state 
objective opt_param

[pharringtonp19]

@duvenaud My understanding may be wrong, but it seems like scan combines two ideas -- (1) the number of times to compose the function and (2) how to compose the function.

Is it possible to construct a function that takes (1) and (2) as separate arguments? From my discussion with @apaszke , here, we wouldn't be able to compile this function because it is a data-dependent function construction. However, we should be able to compile the returned function (if I am understanding things correctly). My sense is that this would be straightforward to implement by defining the function recursively, although since that is not supported at the moment in dex, I am not entirely sure how to go about it.

Example

'### Objective Function 

def objective (x: Float) : (Float & Float)= 
    (log (x*x + 1.0 + sin(3*x)) + 1.5, 1.)

'### Function Composition

def (.*.) (f : Float -> (Float & Float)) (g : Float -> (Float & Float)) : Float -> (Float & Float) =
  
    def g_after_f (x:Float): (Float & Float) = 
        (a1, a2) = f x 
        (b1, b2) = g a1 
        (b1, b2 + a2)

    \(x:Float). g_after_f x

def compose_n_times ... ?

[apaszke]

Yeah I think that's not possible. Think about it this way: if you want to apply f some number n times to an x, you could do it in two ways:

1. You could compose f with itself n times, and then apply to x.
2. You could iterate the application of f to the previous result n times, with initial "previous result" being set to x.

In Haskell both are expressible. The first one is roughly (foldl (.) id (replicate n f)) x, while the latter is fold ($) x (replicate n f). In Dex, only the second one would be accepted, since the first one constructs a data-dependent function definition in the (fold ...) subexpression. In the second expression, the "carry" of the fold is the type of x, which is accepted by Dex for as long as x itself is not a function (or doesn't contain functions as elements). But even in Haskell the second definition is likely to be much faster than the first one.