Skip to content

vllm.utils.jsontree

Helper functions to work with nested JSON structures.

JSONTree module-attribute 

JSONTree = Union[
    dict[str, "JSONTree[_T]"],
    list["JSONTree[_T]"],
    tuple["JSONTree[_T]", ...],
    _T,
]

A nested JSON structure where the leaves need not be JSON-serializable.

_JSONTree module-attribute 

_JSONTree = Union[
    dict[str, "JSONTree[_T]"],
    list["JSONTree[_T]"],
    tuple["JSONTree[_T]", ...],
    dict[str, _T],
    list[_T],
    tuple[_T, ...],
    _T,
]

Same as JSONTree but with additional Union members to satisfy overloads.

_T module-attribute 

_T = TypeVar('_T')

_U module-attribute 

_U = TypeVar('_U')

json_count_leaves 

json_count_leaves(value: JSONTree[_T]) -> int

Count the number of leaves in a nested JSON structure.

Source code in vllm/utils/jsontree.py 
176
177
178
def json_count_leaves(value: JSONTree[_T]) -> int:
    """Count the number of leaves in a nested JSON structure."""
    return sum(1 for _ in json_iter_leaves(value))

json_iter_leaves 

json_iter_leaves(value: JSONTree[_T]) -> Iterable[_T]

Iterate through each leaf in a nested JSON structure.

Source code in vllm/utils/jsontree.py 
39
40
41
42
43
44
45
46
47
48
def json_iter_leaves(value: JSONTree[_T]) -> Iterable[_T]:
    """Iterate through each leaf in a nested JSON structure."""
    if isinstance(value, dict):
        for v in value.values():
            yield from json_iter_leaves(v)
    elif isinstance(value, (list, tuple)):
        for v in value:
            yield from json_iter_leaves(v)
    else:
        yield value

json_map_leaves 

json_map_leaves(
    func: Callable[[Tensor], Tensor],
    value: BatchedTensorInputs,
) -> BatchedTensorInputs

json_map_leaves(
    func: Callable[[_T], _U],
    value: Union[_T, dict[str, _T]],
) -> Union[_U, dict[str, _U]]

json_map_leaves(
    func: Callable[[_T], _U], value: Union[_T, list[_T]]
) -> Union[_U, list[_U]]

json_map_leaves(
    func: Callable[[_T], _U],
    value: Union[_T, tuple[_T, ...]],
) -> Union[_U, tuple[_U, ...]]

json_map_leaves(
    func: Callable[[_T], _U], value: JSONTree[_T]
) -> JSONTree[_U]

json_map_leaves(
    func: Callable[[_T], _U],
    value: Union[BatchedTensorInputs, _JSONTree[_T]],
) -> Union[BatchedTensorInputs, _JSONTree[_U]]

Apply a function to each leaf in a nested JSON structure.

Source code in vllm/utils/jsontree.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def json_map_leaves(
    func: Callable[[_T], _U],
    value: Union["BatchedTensorInputs", _JSONTree[_T]],
) -> Union["BatchedTensorInputs", _JSONTree[_U]]:
    """Apply a function to each leaf in a nested JSON structure."""
    if isinstance(value, dict):
        return {
            k: json_map_leaves(func, v)  # type: ignore[arg-type]
            for k, v in value.items()
        }
    elif isinstance(value, list):
        return [json_map_leaves(func, v) for v in value]
    elif isinstance(value, tuple):
        return tuple(json_map_leaves(func, v) for v in value)
    else:
        return func(value)

json_reduce_leaves 

json_reduce_leaves(
    func: Callable[[_T, _T], _T],
    value: Union[_T, dict[str, _T]],
) -> _T

json_reduce_leaves(
    func: Callable[[_T, _T], _T], value: Union[_T, list[_T]]
) -> _T

json_reduce_leaves(
    func: Callable[[_T" backlink-type="used-by" backlink-anchor="vllm.utils.jsontree.json_reduce_leaves" optional hover>_T, _T], _T],
    value: Union[_T, tuple[_T, ...]],
) -> _T

json_reduce_leaves(
    func: Callable[[_T, _T], _T], value: JSONTree[_T]
) -> _T

json_reduce_leaves(
    func: Callable[[_U, _T], _U],
    value: JSONTree[_T],
    initial: _U,
) -> _U

json_reduce_leaves(
    func: Callable[..., Union[_T, _U]],
    value: _JSONTree[_T],
    initial: _U = cast(_U, ...),
) -> Union[_T, _U]

Apply a function of two arguments cumulatively to each leaf in a nested JSON structure, from left to right, so as to reduce the sequence to a single value.

Source code in vllm/utils/jsontree.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def json_reduce_leaves(
        func: Callable[..., Union[_T, _U]],
        value: _JSONTree[_T],
        initial: _U = cast(_U, ...),  # noqa: B008
        /,
) -> Union[_T, _U]:
    """
    Apply a function of two arguments cumulatively to each leaf in a
    nested JSON structure, from left to right, so as to reduce the
    sequence to a single value.
    """
    if initial is ...:
        return reduce(func, json_iter_leaves(value))  # type: ignore[arg-type]

    return reduce(
        func,  # type: ignore[arg-type]
        json_iter_leaves(value),
        initial,
    )