Support for element-wise repetitions in mlx.core.repeat

[ajithmoola]

Hey,

I've been transitioning my code from JAX to MLX and noticed a potential enhancement that could increase the flexibility of the mlx.core.repeat method. Currently, it seems that mlx.core.repeat only allows for a single integer to specify repetitions for all elements in the input array. However, it would be incredibly useful to support an array-like object that allows specifying a different repetition count for each element in the input array.

a = mlx.core.array([1, 2, 3, 4])
b = mlx.core.array([2, 3, 1, 0])

result = mlx.core.repeat(a, b)
print(result)
# Desired output: array([1, 1, 2, 2, 2, 3])

I can take a stab at this, if this is not on your roadmap.

[awni]

In general MLX does not support operations where the output arrays shape depends on the input arrays data. So the API you are proposing is not so simple to implement. There are some possible alternatives:

• Is it critical that the input containing the repeats is an array or could it be a Python list?
• Do you know the resulting size when you call the repeat? If so we could add a kwarg like Jax (total_repeat_length or something).

Maybe you could say more about your use case to help us understand what would work best.

[ajithmoola]

To give you a better insight into my specific use case, I am currently using jax.numpy.repeat for generating indices to perform scatter-add operations as shown:

num_pts = 100
# segment_lengths is a 1D array of integers with num_pts elements
segment_ids = jnp.repeat(jnp.arange(num_pts), segment_lengths)
n = len(segment_ids)

# prod has the same size as segment_ids along axis=0, prod.shape == (n, 3)
surf = jnp.zeros((num_pts, 3)).at[segment_ids].add(prod)

Since I don't need the output of repeat method to be differentiable, currently I'm using 'numpy.repeat' and converting it to an mlx array as a workaround.

Here's a more comprehensive version of code related to my use case:

import jax.numpy as jnp
import numpy as np


def prepare_data_for_evaluation_jax(
    PHI, ac_spans, num_supp, ctrl_pts, ac_cells_ac_supp, fn_sh
):
    max_lev = max(ctrl_pts.keys())
    nCP = np.zeros(max_lev + 2, dtype=np.int_)
    CP_dim = ctrl_pts[0].shape[-1]

    segment_lengths = num_supp
    num_pts = segment_lengths.size
    segment_ids = jnp.repeat(jnp.arange(num_pts), segment_lengths)

    PHI = PHI.astype(jnp.float32).reshape(-1, 1)

    for lev in range(1, max_lev + 2):
        nCP[lev] = nCP[lev - 1] + np.prod(fn_sh[lev - 1])

    ctrl_pts_flat = [
        jnp.array(ctrl_pts[lev]).reshape(-1, CP_dim).astype(jnp.float32)
        for lev in range(max_lev + 1)
    ]
    ctrl_pts = jnp.vstack(ctrl_pts_flat)

    Jm = [
        nCP[fn_lev] + np.ravel_multi_index(fnIdx, fn_sh[fn_lev])
        for cell_lev, cellIdx in ac_spans
        for fn_lev, fnIdx in ac_cells_ac_supp[cell_lev][cellIdx]
    ]

    Jm = jnp.array(Jm)

    prod = PHI * ctrl_pts[Jm]
    output = jnp.zeros((num_pts, ctrl_pts.shape[1])).at[segment_ids].add(prod)
    return output

[awni]

Thanks for the details.

If the segment_lengths is a static array in you computation (e.g. its an input numpy array or list or something like that) its not inconceivable that we could support this version of repeat. That said, if its static It also might just make sense to take the repeat part out of the graph and make it in numpy (as you are doing now for the workaround). I'm not sure if it changes in your computation from iteration to iteration, but if not that might be more efficient as well.

As a next step, I'd file an issue if this is useful for you. We can mark it as an enhancement and then prioritize from there.

[ajithmoola]

Thank you for your feedback and the proposed next steps. In my current application, segment_lengths does not change from iteration to iteration. I think the current implementation with numpy is good enough for this particular use case.

I have one more query regarding optimizers in mlx. Should I create a new discussion or I can continue here?

[awni]

A new discussion with an appropriate title is most useful for discovery purposes (so others are more likely to benefit from the question). Thanks!