Make `zeros_like` work with shapeless compile / export.

[josharian]

Hi! Thanks very much for mlx.

I'm trying to use shapeless exports and I'm hitting "cannot infer output shapes" for a few ops that seem like they should support shapeless use.

The first one I hit was Conv. I tried adding Conv::output_shape, and it appears to work. Should I send a PR with this, or is there some reason it isn't implemented?

The next was zeros_like. It uses Full, which doesn't implement output_shape. But zeros_like seems naively like it should be sufficiently shapeless. Should there be a FullLike? Or some other more direct way to do zeros_like and ones_like in a shapeless manner? Or am I missing something here?

Thanks!

[awni]

For the convolution, yes feel free to send the PR.

zeros_like needs some thought. For example suppose you export / compile a function like:

def fun(x)
  return zeros_like(x)

That should infer the output shape from the input shape. But zeros_like just calls full(shape, 0). And that doesn't preserve the fact that x was an input in the graph which will break shape inference.

[josharian]

Conv: Great, thanks! Will send the PR tonight or tomorrow.

zeros_like: That makes sense. I haven't dug much into the mlx internals, but maybe could there be a full_like, so that zeros_like calls through to full_like(x, 0)? Then Full could be straightforwardly implemented in terms of FullLike.

[awni]

zeros_like: That makes sense. I haven't dug much into the mlx internals, but maybe could there be a full_like, so that zeros_like calls through to full_like(x, 0)? Then Full could be straightforwardly implemented in terms of FullLike.

Yea something like that could work to support shapeless zeros_like. I think we can do it without adding a new primitive. Basically:

• zeros_like and ones_like call full_like
• full_like does a broadcast_arrays (which is safe for shapeless)
• Then the full_like calls a new full_impl implementation which does not broadcast

The original zeros, ones call full which does a broadcast_to and then calls the new full implementation helper.