FR: expose device/sharding argument for NamedArray.wrap

[amifalk]

No description provided.

[danieldjohnson]

Hm, I'm not sure this is necessary? You can provide the device and sharding for the array before wrapping it, and pz.nx.wrap should preserve it, e.g.

arr = pz.nx.wrap(jax.device_put(my_array, my_sharding))

[amifalk]

Yep, just a little syntactic sugar :). When training large models, I often use a data loader to grab data from the disk and process it onto a numpy array, then directly from there convert it to a NamedArray sharded over devices.