Skip to content

vllm.utils.jsontree

Helper functions to work with nested JSON structures.

Functions:

Attributes:

  • JSONTree (TypeAlias) –

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

JSONTree = dict[str, 'JSONTree[_T]'] | list['JSONTree[_T]'] | tuple['JSONTree[_T]', ...] | _T module-attribute

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

_JSONTree = dict[str, 'JSONTree[_T]'] | list['JSONTree[_T]'] | tuple['JSONTree[_T]', ...] | dict[str, _T] | list[_T] | tuple[_T, ...] | _T module-attribute

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

json_count_leaves(value)

Count the number of leaves in a nested JSON structure.

Source code in vllm/utils/jsontree.py 
155
156
157
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(value)

Iterate through each leaf in a nested JSON structure.

Source code in vllm/utils/jsontree.py 
36
37
38
39
40
41
42
43
44
45
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(func, value) 

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

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

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

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

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

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

Source code in vllm/utils/jsontree.py 
83
84
85
86
87
88
89
90
91
92
93
94
95
def json_map_leaves(
    func: Callable[[_T], _U],
    value: Any,
) -> "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) for k, v in value.items()}  # type: ignore
    elif isinstance(value, list):
        return [json_map_leaves(func, v) for v in value]  # type: ignore
    elif isinstance(value, tuple):
        return tuple(json_map_leaves(func, v) for v in value)
    else:
        return func(value)

json_reduce_leaves(func, value, initial=...) 

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

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

json_reduce_leaves(
    func: Callable[[_T, _T], _T], value: _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

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 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def json_reduce_leaves(
    func: Callable[[_T, _T], _T] | Callable[[_U, _T], _U],
    value: _JSONTree[_T],
    initial: _U = ...,  # type: ignore[assignment]
    /,
) -> _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

    return reduce(func, json_iter_leaves(value), initial)  # type: ignore