Skip to content

laplax.util.ops

Contains operations for flexible/adaptive compute.

str_to_bool 

str_to_bool(value: str) -> bool

Convert a string representation of a boolean to a boolean value.

Parameters:

Name Type Description Default
value str

A string representation of a boolean ("True" or "False").

 required

Returns:

Type Description
bool

The corresponding boolean value.

Raises:

Type Description
ValueError

If the string does not represent a valid boolean value.
Source code in laplax/util/ops.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def str_to_bool(value: str) -> bool:
    """Convert a string representation of a boolean to a boolean value.

    Args:
        value: A string representation of a boolean ("True" or "False").

    Returns:
        The corresponding boolean value.

    Raises:
        ValueError: If the string does not represent a valid boolean value.
    """
    valid_values = {"True": True, "False": False}
    if value not in valid_values:
        msg = "invalid string representation of a boolean value"
        raise ValueError(msg)
    return valid_values[value]

precompute_list 

precompute_list(func: Callable, items: Iterable, precompute: bool | None = None, **kwargs: Kwargs) -> Callable

Precompute results for a list of items or return the original function.

If option is enabled, this function applies func to all items in items and stores the results for later retrieval. Otherwise, it returns func as-is.

Parameters:

Name Type Description Default
func Callable

The function to apply to each item in the list.

 required
items Iterable

An iterable of items to process.

 required
precompute bool | None

Determines whether to precompute results: - None: Use the default precompute setting. - bool: Specify directly whether to precompute.

 None
**kwargs Kwargs

Additional keyword arguments, including: - precompute_list_batch_size: Batch size for precomputing results.

 {}

Returns:

Type Description
Callable

A function to retrieve precomputed elements by index, or the original func if precomputation is disabled.
Source code in laplax/util/ops.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def precompute_list(
    func: Callable, items: Iterable, precompute: bool | None = None, **kwargs: Kwargs
) -> Callable:
    """Precompute results for a list of items or return the original function.

    If `option` is enabled, this function applies `func` to all items in `items`
    and stores the results for later retrieval. Otherwise, it returns `func` as-is.

    Args:
        func: The function to apply to each item in the list.
        items: An iterable of items to process.
        precompute: Determines whether to precompute results:
            - None: Use the default precompute setting.
            - bool: Specify directly whether to precompute.
        **kwargs: Additional keyword arguments, including:
            - precompute_list_batch_size: Batch size for precomputing results.

    Returns:
        A function to retrieve precomputed elements by index, or the original
            `func` if precomputation is disabled.
    """
    if precompute is None:
        precompute = DEFAULT_PRECOMPUTE_LIST

    if precompute:
        precomputed = jax.lax.map(
            func, items, batch_size=kwargs.get("precompute_list_batch_size")
        )

        def get_element(index: int):
            return jax.tree.map(operator.itemgetter(index), precomputed)

        return get_element

    return func