Imported Functions called with an unknown dimension size.

[cianciosa]

I can define a function and compile then call it with a variable number in the leading dimension.

import mlx.core

input_shape = 5
output_shape = 10

input_buffer = mlx.core.ones((input_shape))
batch_input = mlx.core.ones((10000, input_shape))

weight = mlx.core.ones((output_shape,input_shape))
@mlx.core.compile
def func(x):
    return mlx.core.einsum('ij,...j->...i', weight, x)

print(func(input_buffer))
print(func(batch_input))

Running this, the first two functions return

array([5, 5, 5, ..., 5, 5, 5], dtype=float32)
array([[5, 5, 5, ..., 5, 5, 5],
       [5, 5, 5, ..., 5, 5, 5],
       [5, 5, 5, ..., 5, 5, 5],
       ...,
       [5, 5, 5, ..., 5, 5, 5],
       [5, 5, 5, ..., 5, 5, 5],
       [5, 5, 5, ..., 5, 5, 5]], dtype=float32)

I'm now trying to export the function. This now requires requires specifying an input.

mlx.core.export_function("test.mlxfn", func, input_buffer)

import_fn = mlx.core.import_function("test.mlxfn")

print(import_fn(input_buffer))
print(import_fn(batch_input)) # This fails because it has the wrong input dimension.

This fails when I call it with my batch_input

(array([5, 5, 5, ..., 5, 5, 5], dtype=float32),)
Traceback (most recent call last):
  File "/Users/m4c/Desktop/mlx_example/test.py", line 49, in <module>
    print(import_fn(batch_input))
          ^^^^^^^^^^^^^^^^^^^^^^
ValueError: [import_function::call] No imported function found which matches the given positional and keyword arguments. Possible functions include:
 1. Function with 1 positional inputs and 0 keyword inputs:
   1: (5) float32

Called with 1 positional inputs and 0 keyword inputs:
  1: (10000,5) float32

Is there away to call this function so I use a variable batch size? When I tried using shapeless=True

mlx.core.export_function("test.mlxfn", func, input_buffer, shapeless=True)

import_fn = mlx.core.import_function("test.mlxfn")

print(import_fn(input_buffer))
print(import_fn(batch_input))

This fails with the output

Traceback (most recent call last):
  File "/Users/m4c/Desktop/mlx_example/test.py", line 48, in <module>
    print(import_fn(input_buffer))
          ^^^^^^^^^^^^^^^^^^^^^^^
ValueError: [Primitive::output_shapes] Full cannot infer output shapes.

[angeloskath]

Shapeless export doesn't allow the number of dimensions or the datatype to change. So if you change

input_buffer = mlx.core.ones((input_shape))

to

input_buffer = mlx.core.ones((1, input_shape))

and export with shapeless=True it should work (at least it works for me).

You can also export several functions in the same file with

with mx.exporter("file.mlxfn", func) as exporter:
    exporter(my_inputs)
    exporter(my_other_inputs)

[cianciosa]

I tried

input_buffer = mlx.core.ones((1, input_shape))

but it still gives me the same error. But after updating my mlx install it corrected.