I'm getting strange behavior when trying to use compile.

[cianciosa]

A apologize for the size of the example but this is smallest I could make this example and still trigger the issue.

#include <iostream>
#include <numbers>

#include <mlx/mlx.h>

typedef std::vector<mlx::core::array> outputs;
typedef const outputs &inputs;
typedef std::function<outputs(inputs)> function;

class fourier_context {
private:
    const mlx::core::array u_angles;

    static mlx::core::array make_u_angles() {
        return mlx::core::linspace(0, std::numbers::pi_v<float>, 10,
                                   mlx::core::float32);
    }

    static mlx::core::array
    make_cosuv(const mlx::core::array u_angles) {
        const mlx::core::array xmu = mlx::core::einsum("i,j->ij", {
            mlx::core::ones({6}, mlx::core::int32),
            u_angles
        });
        return mlx::core::cos(xmu) + mlx::core::sin(xmu);
    }

    static mlx::core::array make_norm(const mlx::core::array u_angles,
                                      const mlx::core::array parity) {
        const mlx::core::array dnorm = mlx::core::ones_like(u_angles);

        const mlx::core::array mask =
            mlx::core::where(u_angles == 0 ||
                             u_angles == std::numbers::pi_v<float>,
                             dnorm, dnorm);
        return mlx::core::einsum("ij,j->ij", {parity, mask});
    }

public:
    const mlx::core::array cosuv;
    const mlx::core::array invcosuv;

    fourier_context() :
    u_angles(make_u_angles()),
    cosuv(make_cosuv(u_angles)),
    invcosuv(make_norm(u_angles, cosuv)) {}

    function compile_to_real(const mlx::core::array trig_parity) const {
        return mlx::core::compile([trig_parity](inputs in) -> outputs {
            return {mlx::core::einsum("...i,ij->...j",
                                      {in[0], trig_parity})};
        });
    }

    mlx::core::array to_real(const mlx::core::array amn,
                             const mlx::core::array trig_parity) const {
        return compile_to_real(trig_parity)({amn})[0];
    }

    function compile_to_fourier(const mlx::core::array trig_parity) const {
        return mlx::core::compile([trig_parity](inputs in) -> outputs {
            return {mlx::core::einsum("...j,ij->...i",
                                      {in[0], trig_parity})};
        });
    }

    mlx::core::array to_fourier(const mlx::core::array aij,
                                const mlx::core::array trig_parity) const {
        return compile_to_fourier(trig_parity)({aij})[0];
    }
};

int main(int argc, const char * argv[]) {
    const fourier_context fc;

    const mlx::core::array amn = mlx::core::ones({6}, mlx::core::float32);
    const mlx::core::array amnc = fc.to_fourier(fc.to_real(amn,
                                                           fc.cosuv),
                                                fc.invcosuv);

    std::cout << amnc << std::endl;
    return 0;
}

When I run this, I get an error when it gets to the line std::cout << amnc << std::endl;

Compiler failed to build request
libc++abi: terminating due to uncaught exception of type std::runtime_error: [metal::Device] Unable to build metal library from source
program_source:1485:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1511:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1539:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1568:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1595:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1622:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1649:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1676:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1711:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1749:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1787:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1825:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1863:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1901:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1939:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:1977:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:2016:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^
program_source:2055:28: error: use of undeclared identifier 'tmp_K'
  float tmp_J = Subtract()(tmp_K, tmp_I);
                           ^

I can make this crash go away if I do any of these three actions.

• If I remove the where op in make_norm.
• If I replace the u_angles in make_cosuv with mlx::core::ones_like(u_angles).
• If I remove the mlx::core::compile functions and evaluate it directly.

[awni]

There is definitely a bug here.. as it should never fail to compile like that. I believe the bug is related to the fact that you have the same variable as an implicit capture to two different functions you are compiling which is causing the graph to get in a bad state. We should fix it.

However, I would strongly advise against compiling with captured variables like that. Here's a quick explanation:

a = constant;
b = constant;
x = a + b;
auto fun = [x](array y) {
   return y + x;
}

In that case MLX needs to compile your function but it doesn't know that x is an implicit input. So essentially it will try to find the entire graph and compile it. And so what MLX is really compiling will be this:

auto fun = [a, b](array y) {
   return y + (a + b);
};

which is usually not what you want. So instead you should just make x an argument to the lambda as well:

a = constant;
b = constant;
x = a + b;
auto fun = [](array y, array x) {
   return y + x;
}

[awni]

In addition to the above, also avoid compiling lambdas as it's very easy to trigger a recompile every time you call that function.

In your case every time you call compile_to_fourrier it will recompile the lambda since it's making a new object there. Instead you can persist the compiled function (in like a std::funciton or something) or just compile a free function which won't get recompiled since it's not changing.

[cianciosa]

It makes sense to evaluate anything that's going to be constant inside the lambda before capture. The thing that motivated this pattern was I was trying to build a sequential Keras model in a loop without having every layer be an input to the final function, or having a bunch of non-mlx code inside the lambda.

The basic idea was (reusing my typedefs for brevity)

// This layer acts as a placeholder.
function func = [](inputs x) -> outputs {
     return x;
 };

 for (size_t i = 0; i < 5; i++) {
// Non-mlx code to load weights and biases.
    mlx::core::array a = mlx::core::ones({1});

// Create a new layer from the previous layer.
    func = [a, func](inputs x) {
        return {a + func(x)[0]};
    };
}

std::cout << func({mlx::core::zeros({1})})[0] << std::endl;

So this motivates a few questions. For non-compiled functions, anything that's not mlx code in a lambda will get run every time I call that lambda.

What happens in a mlx compiled function containing non-mlx routines?

I ask this because I can see an alternative way todo this like

function func = [](inputs x) -> outputs {
    for (size_t i = 0; i < 5; i++) {
        mlx::core::array a = non_mlx_code_to_load_array_from_disk();
        return {x[0] = a + x[0]};
    }
};

It's clear if I called this without compiling it, it would reload a from disk every time which you wouldn't want. Does a compiled function still run that loop every time it's called, or does the mlx compiler run the lambda and capture every mlx operation so the non-mlx code is not evaluated?

As an aside, while trying to prepare for this response, I learned something new about python. So ignoring all mlx code and just using plan C++ and Python, my first pattern doesn't even work in python.

The C++ code

#include <iostream>
#include <functional>

int main() {
    std::function<int(int)> func = [](int x) {
        return x;
    };

    for (size_t i = 0; i < 5; i++) {
        int a = 1;
        func = [a, func](int x) {
            return a + func(x);
        };
    }

    std::cout << func(0) << std::endl;

    return 0;
}

Outputs 5 as I would expect. The closest python equivalent (Although this maybe wrong. I'm not an expert in python)

func = lambda x : x

for i in range(5):
    a = 1
    
    func = lambda x: a + func(x)

print(func(0))

Fails to run resulting in an infinite loop.

python3 test.py
Traceback (most recent call last):
  File "test.py", line 8, in <module>
    print(func(0))
          ^^^^^^^
  File ".../test.py", line 6, in <lambda>
    func = lambda x: a + func(x)
                         ^^^^^^^
  File ".../test.py", line 6, in <lambda>
    func = lambda x: a + func(x)
                         ^^^^^^^
  File ".../test.py", line 6, in <lambda>
    func = lambda x: a + func(x)
                         ^^^^^^^
  [Previous line repeated 996 more times]

[awni]

For reference, here is a simple python reproduction of the bug:

import mlx.core as mx                                                          
    
t = mx.ones((10,))                                                             
u = (1.0 - t) * 0.0 + t * 3.0
  
o = mx.ones((6,))                                                              
b = o[:, None] * u
  
c = b * mx.ones_like(u)                                                        

a = mx.ones((6,))
d = mx.compile(lambda x: x @ b)(a)                                             
e = mx.compile(lambda x: x @ c.T)(d)
mx.eval(e)

It can be fixed by either evluating b and c before compile or by making them explicit inputs to the compiled lambdas.