def wrapped(*args, **kwargs):
    if kwargs:
      raise NotImplementedError("sharded_jit over kwargs not yet supported")

    f = lu.wrap_init(fun)
    if static_argnums:
      if max(static_argnums) >= len(args):
        raise ValueError(
            f"jitted function has static_argnums={static_argnums}"
            f" but was called with only {len(args)} positional "
            f"argument{'s' if len(args) > 1 else ''}. "
            "All static broadcasted arguments must be passed positionally.")
      dyn_argnums = [i for i in range(len(args)) if i not in static_argnums]
      f, args = argnums_partial(f, dyn_argnums, args)

    args_flat, in_tree = tree_flatten((args, kwargs))
    in_parts_flat = tuple(flatten_axes("sharded_jit in_parts",
                                       in_tree.children()[0], in_parts))
    if local_in_parts is not None:
      local_in_parts_flat = tuple(flatten_axes("sharded_jit local_in_parts",
                                               in_tree.children()[0], local_in_parts))
    else:
      local_in_parts_flat = None

    flat_fun, out_tree = flatten_fun(f, in_tree)
    # TODO(skye): having a function-typed param in a primitive seems dicey, is
    # there a better way?
    out_parts_thunk = HashableFunction(
        lambda: tuple(flatten_axes("sharded_jit out_parts", out_tree(), out_parts)),
        closure=out_parts)
    if local_out_parts:
      local_out_parts_thunk = HashableFunction(
          lambda: tuple(flatten_axes("sharded_jit local_out_parts",
                                     out_tree(), local_out_parts)),
          closure=local_out_parts)
    else:
      local_out_parts_thunk = HashableFunction(lambda: None, closure=None)

    out = sharded_call(
        flat_fun,
        *args_flat,
        nparts=nparts,
        in_parts=in_parts_flat,
        out_parts_thunk=out_parts_thunk,
        local_in_parts=local_in_parts_flat,
        local_out_parts_thunk=local_out_parts_thunk,
        local_nparts=local_nparts,
        name=flat_fun.__name__)
    return tree_unflatten(out_tree(), out)