Unable to reproduce the example

[arn4]

I was trying to reproduce the code in the example of the README file, but it crashes once the request is sent.

In particular, on the second machine I get this error message:

Received tensor request
Error processing tensor: [reshape] Cannot reshape array of size 163840 into shape (1,40,8192).

After a quick look at the code, it seems to me that the function Byte_to_tensor is outputting a mx.array twice as long as it should be, but I was unable to find the error.

Do you have any idea why this is happening?

I've also tried with different models, but the same error is happining (the numbers are different, but the tensor that needs to be reshaped is always long the double as it should be)

[mzbac]

It looks like there are some layer mismatch somehow. Would you be able to install the mlx-sharding and try the example at https://github.com/mzbac/mlx_sharding?tab=readme-ov-file#running-the-servers ?

[arn4]

Ok, using the first 2 commands it works.

I've tested both the model of the example and a larger one.