mx.compile for class methods

[romanoneg]

Is there a way to compile a class method through mx.compile?

I can't seem to find how to do it while I know that in jax you can do it in a couple different ways.

The helper function way shown in the jax docs linked above might work but it feels like there should be a different way of doing it using mx.compile's input and output.

some example code:

class Test:
    @mx.compile
    def method(self, x):
        return x
t = Test()
print(t.method(1)) # TypeError: Test.method() missing 1 required positional argument: 'x'

the 1 here gets passed in as self into the method, and trying to do t.method(t, 1) leads to hashing errors.

Have I missed a very obvious fix to this issue?

[awni]

You can do it like this:

import mlx.core as mx

class Test:
    @mx.compile
    def method(self, x):
        return x

t = Test()
print(mx.compile(t.method)(mx.array(1.0)))

Note though that the method capture self there so it might not be want you want.

What you have doesn' t work, because compile currently doesn't deal with instance methods (so the compiled function capture t). You could almost get it to work with :

t.method(t, 1)

But then t is not a type we can compile with. But you could make it one by inheriting from say dict (if it makes sense for you:

import mlx.core as mx

class Test(dict):

    @mx.compile
    def method(self, x):
        return x

t = Test()
print(t.method(t, mx.array(1.0)))

I don't think that's a particularly good idea to do though if it can be avoided. Better to stick to compiling free functions for now.