How to make jax.numpy.linspace to produce sharded result

[jczaja]

Hi,

How can I make jax.numpy.linspace to produce result already sharded across my two jax devices?

Context:

I have a host with two JAX devices. Problem is that each of this devices are limited in memory. So I can call jax.numpy.linspace(start, end, num=50000000) but bigger values of num would call Out of memory error. So if it is possible
to make jax.numpy.linspace to directly generate numbers on both of my devices then I could have jax.numpy.linspace(start,end,num=2*50000000) working. Is this possible?

[yashk2810]

One way of doing this is like this: (the example is with arange but you can adopt it for linspace)

@jax.jit
def f():
  x = jnp.arange(16).reshape(8, 2)
  return jax.lax.with_sharding_constraint(x, sharding)

I assume you know what the sharding is?

[Findus23]

Thank you for the answer as I was wondering something similar recently.
Just to make sure, using out_shardings on the function would also do the same thing here:

def f():
    return jax.numpy.arange(16).reshape(8, 2)

f = jax.jit(f, out_shardings=sharding)

[jczaja]

Thanks very much for answers