Skip to content

Operations

fori_loop(lower: int, upper: int, f: Callable[[int], PyTree | None]) -> PyTree | None

Loop from lower to upper evaluating a function f.

Parameters:

Name Type Description Default
lower int

The starting index (inclusive).

 required
upper int

The ending index (exclusive).

 required
f Callable[[int], PyTree | None]

A function accepting an integer index i.

 required

Returns:

Type Description
PyTree | None

A PyTree whose leaves are stacked with the evaluations of f (which reduces to None if nothing is returned).
Source code in src/bayinx/ops.py 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def fori_loop(
    lower: int,
    upper: int,
    f: Callable[[int], PyTree | None]
) -> PyTree | None:
    """
    Loop from `lower` to `upper` evaluating a function `f`.

    Parameters:
        lower: The starting index (inclusive).
        upper: The ending index (exclusive).
        f: A function accepting an integer index `i`.

    Returns:
        A `PyTree` whose leaves are stacked with the evaluations of `f` (which reduces to `None` if nothing is returned).
    """
    from bayinx.core.context import _model_context

    # Create the sequence of indices to iterate over
    idxs = jnp.arange(lower, upper)

    # Check for existing model context
    within_context = hasattr(_model_context, "target")

    if within_context:
        # Reference the outer model context
        outer_target = _model_context.target

    # Scan to accumulate log-probabilities
    def scanner(carry, i):
        # Shadow the outer model context in the local scope
        if within_context:
            _model_context.target = Target(jnp.array(0.0))

        try:
            # Evaluate user callable with loop index
            user_out = f(i)

            if within_context:
                # Extract the local target density
                local_lp = _model_context.target.value
            else:
                local_lp = 0.0

        finally:
            if within_context:
                # Restore model context to outer scope
                _model_context.target = outer_target

        if within_context:
            # Accumulate local target density into carry
            carry = carry + local_lp

        return carry, user_out

    # Run scan
    total_log_prob, user_results = lax.scan(scanner, 0.0, idxs)

    if within_context:
        # Update outer target with accumulated local target from the loop
        outer_target += total_log_prob

    # Return unwrapped none when there are no user outputs
    if all(x is None for x in jt.flatten(user_results)):
        return None

    return user_results

map(f: Callable[..., PyTree | None], *xs: Any) -> PyTree | None

Map a function over the leading axis of the arguments.

Parameters:

Name Type Description Default
f Callable[..., PyTree | None]

A user-defined function that accepts slices of the input positional arguments.

 required
xs Any

Additional positional arguments that are sliced and passed to f.

 ()

Returns:

Type Description
PyTree | None

A PyTree whose leaves are stacked with the evaluations of f (which reduces to None if nothing is returned).
Source code in src/bayinx/ops.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
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
def map(
    f: Callable[..., PyTree | None],
    *xs: Any
) -> PyTree | None:
    """
    Map a function over the leading axis of the arguments.

    Parameters:
        f: A user-defined function that accepts slices of the input positional arguments.
        xs: Additional positional arguments that are sliced and passed to `f`.

    Returns:
        A `PyTree` whose leaves are stacked with the evaluations of `f` (which reduces to `None` if nothing is returned).
    """
    from bayinx.core.context import _model_context

    # Check for existing model context
    within_context = hasattr(_model_context, "target")

    if within_context:
        # Reference the outer model context
        outer_target = _model_context.target

    # Scan to accumulate log-probabilities
    def scanner(carry, x):
        if within_context:
            # Shadow the outer model context in the local scope
            _model_context.target = Target(jnp.array(0.0))

        try:
            # Evaluate user callable (which uses the local scope's model context)
            user_out = f(*x)

            if within_context:
                # Extract the local target density
                local_lp = _model_context.target.value
        finally:
            if within_context:
                # Restore model context to outer scope
                _model_context.target = outer_target

        if within_context:
            # Accumulate local target density
            carry = carry + local_lp

        return carry, user_out

    # Run scan
    total_log_prob, user_results = lax.scan(scanner, 0.0, xs)

    if within_context:
        # Update outer target with accumulated local target
        outer_target += total_log_prob

    # Return unwrapped none when there are no user outputs
    if all(x is None for x in jt.flatten(user_results)):
        return None

    return user_results