def sharded_jit(
    fun: Callable,
    in_parts,
    out_parts,
    num_partitions: Optional[int] = None,
    local_in_parts=None,
    local_out_parts=None,
    local_num_partitions=None,
    static_argnums: Union[int, Iterable[int]] = (),