Query: Do we need to partition the weights also? Or jax.jit will take of it?

[MohitKumar020291]

From this tutorial
there is code

class Params(NamedTuple):
    w1: jnp.ndarray
    w2: jnp.ndarray
    
@jit
def ffn(x: jnp.array, params: Params):
    z = jnp.maximum(x @ params.w1, 0)
    return z @ params.w2

• if the x is sharded but the weights (w1) are not - would not that be a problem or jit will figure it out because at documentation states
  • "Once you have sharded data, the easiest way to do parallel computation is to simply pass the data to a jax.jit()-compiled function! In JAX, you need to only specify how you want the input and output of your code to be partitioned, and the compiler will figure out how to: 1) partition everything inside; and 2) compile inter-device communications."
  • it is contradictory to me whether jax.jit wants a sharded weights or it will take care of it

PS: Cannot try on my own because of limited resources.