Skip to content

[FEATURE] mx.repeat with variable number of repetitions #1785

@kyrollosyanny

Description

@kyrollosyanny

Describe the bug
In numpy, np.repeat can accept both int and array type for the repetition. The main use case for passing an array is if each element in the original array needs to repeated a different number of times. Currently, this fails in MLX.

To Reproduce
In numpy:

a=np.array([1,2,3,4])

print(np.repeat(a,a))

[1 2 2 3 3 3 4 4 4 4]

in MlX, this does not work since repetition needs to be an integer.

It would be great if we can add this behavior to mx.repeat unless there is another efficient way to do this.

Thanks

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions