Did you handle replicate for 1 device?

[jiapei-nexera]

Looks like 1 device is NOT replicatable, but NOT asserted or let it be and successful ...

See Also:
    - device_put
    - device_put_sharded
  """
  if not isinstance(devices, Sequence) or not devices:
    raise ValueError("`devices` argument to `device_put_replicated must be "
                     "a non-empty sequence.")
  def _device_put_replicated(x):
    aval = core.unmapped_aval(len(devices), core.no_axis_name, 0,
                              core.raise_to_shaped(core.get_aval(x)))
    assert (isinstance(aval, ShapedArray) and
            len(xla.aval_to_xla_shapes(aval)) == 1)
    buf, = dispatch.device_put(x, devices[0])
    rest_bufs = [buf.copy_to_device(d) for d in devices[1:]]
    if config.jax_array:
      from jax.experimental import array, sharding
      sharding_spec = pxla._create_pmap_sharding_spec(aval)
      return array.Array(
          aval.shape, sharding.PmapSharding(np.array(devices), sharding_spec),
          [buf, *rest_bufs], committed=True)
    else:
      return pxla.make_sharded_device_array(aval, None, [buf, *rest_bufs])

  with config_explicit_device_put_scope():
    return tree_map(_device_put_replicated, x)