Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Inconsistency in compile with kwargs #825

Open
awni opened this issue Mar 12, 2024 · 2 comments
Open

Inconsistency in compile with kwargs #825

awni opened this issue Mar 12, 2024 · 2 comments

Comments

@awni
Copy link
Member

awni commented Mar 12, 2024

The fact that some of the following work but some don't seems inconsistent and unexpected. Filling this here mostly so I don't forget about it.

import mlx.core as mx

@mx.compile
def fun(x, y=None):
    if y is not None:
        return x + y
    else:
        return x + 1

fun(mx.array(1.0)) # ok
fun(mx.array(1.0), mx.array(2.0)) # ok
fun(mx.array(1.0), None) # exception
fun(mx.array(1.0), y=None) #exception
@romanoneg
Copy link

Second time writing this comment sorry for any notification spam, and let me know if this is the wrong place for it, it seems related enough, but if its the wrong place let me know and I'll delete.

I've run into another edge case with mx.compile and custom dataclasses and Im not super certain why the behavior is occurring:

import mlx.core as mx
from collections import namedtuple

exampleClass = namedtuple('Example', ['x','y'])
example_tuple = exampleClass(x=0,y=1)

def foo(mytuple):
    return mytuple[0] + mytuple[1] 

print(foo(mx.array([0,1]))) # works outputs array(1, dtype=int32)
print(foo(example_tuple)) # works outputs 1

compiled_foo = mx.compile(foo)

print(compiled_foo(mx.array([0,1]))) # outputs array(1, dtype=int32)
print(compiled_foo(example_tuple)) # outputs None (?huh?)

@awni
Copy link
Member Author

awni commented Mar 15, 2024

In the named example type you are not doing any array options so compiling through that doesn't make sense. (The 0 and 1 never get cast to mx.array. You can fix it by doing:

exampleClass = namedtuple('Example', ['x','y'])
example_tuple = exampleClass(x=mx.array(0),y=mx.array(1))

def foo(mytuple):
    return mytuple[0] + mytuple[1]

compiled_foo = mx.compile(foo)
print(compiled_foo(example_tuple))

We should have better error messaging (or find a way to support it).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

2 participants