How does mx.pad work?

[sck-at-ucy]

I am trying to use mx.pad with the optional constant array containing the padding values to be inserted, but it seems that only the first element of the constant array is utilized whatever I try. Am I interpreting the description of mlx.core.pad() the wrong way?

`import numpy as np import mlx.core as mx

nx, ny = 10, 10 # Set grid dimensions

x = mx.array(np.linspace(0,1,nx)) y = mx.array(np.linspace(0,1,ny))

T = mx.zeros([nx, ny])

Tconst = mx.array(np.linspace(9,0,10))

print(Tconst)

T = mx.pad(T, ((1, 1), (1, 1)),Tconst)

print(T)`

array([9, 8, 7, ..., 2, 1, 0], dtype=float32) array([[9, 9, 9, ..., 9, 9, 9], [9, 0, 0, ..., 0, 0, 9], [9, 0, 0, ..., 0, 0, 9], ..., [9, 0, 0, ..., 0, 0, 9], [9, 0, 0, ..., 0, 0, 9], [9, 9, 9, ..., 9, 9, 9]], dtype=float32)

Process finished with exit code 0

[machineko]

Padding with constant in almost every framework use scalar value and its same here (I feel like passing array is some sort of performance optimisation (?) or maybe extra usability cause original code still use only first array element) You can check original code -> https://github.com/ml-explore/mlx/blob/430bfb494431fbe36c3f635bd7296a356dd772e6/mlx/ops.cpp#L571C1-L577C34

[sck-at-ucy]

Thanks for pointing me to the original code, I think you're right, I had hopes that the reference to an array was hinting at something more exciting :)

[awni]

Good question, right now the constant_values is really just a scalar.

We should probably throw or something if the wrong shape is provided.

I think in Numpy the constant_values need to be broadcastable to the shape of the pad_width. So if pad_width is ((a, b), (c, d))then theconstant_valuesneeds to broadcast to a shape of[2, 2]`. We can add this as a feature in the future if it's useful.

Also (unrelated) it looks like we should add a linspace :). If that's pretty useful for you, feel free to file another issue.

[sck-at-ucy]

Thank you this clarifies. Yes, MLX native linspace would be handy to have.