Skip to content

Core Utilities Documentation

GPU Utilities

_block_until_ready 

_block_until_ready(value)

Ensures GPU is done doing work on input before returning.

Source code in dlem/core.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
def _block_until_ready(value):
    """
    Ensures GPU is done doing work on input before returning.

    """
    if isinstance(value, np.ndarray):
        return
    if hasattr(value, "block_until_ready"):
        value.block_until_ready()
        return
    if isinstance(value, Mapping):
        for v in value.values():
            _block_until_ready(v)
        return
    if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
        for v in value:
            _block_until_ready(v)

Data Utilities

normalize_expected_observed 

normalize_expected_observed(band, log=True)

Compute correlation and MSE between generated and target bands in the expected/observed normalised space. Used for tracking training progress.

Source code in dlem/core.py 
 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
def normalize_expected_observed(band, log=True):
    """
    Compute correlation and MSE between generated and 
    target bands in the expected/observed normalised space.
    Used for tracking training progress.

    """    
    band = jnp.asarray(band, dtype=jnp.float32)
    if band.ndim != 2:
        raise ValueError("normalize_expected_observed expects a 2D array.")

    width, n = band.shape
    if width == 0 or n == 0:
        return jnp.zeros_like(band)

    row_idx = jnp.arange(width, dtype=jnp.int32)[:, None]
    col_idx = jnp.arange(n, dtype=jnp.int32)[None, :]
    mask = col_idx >= row_idx

    masked_band = jnp.where(mask, band, 0.0)
    counts = mask.sum(axis=1).astype(band.dtype)
    sums = masked_band.sum(axis=1)
    means = jnp.where(counts > 0, sums / counts, 0.0)
    means_expanded = means[:, None]
    offset = 1e-6
    valid = mask & (means_expanded > 0.0)
    denom = jnp.where(valid, means_expanded, 1.0)
    ratios = jnp.where(valid, (band + offset) / (denom + offset), 1.0)
    if not log:
        normalized = jnp.where(valid, ratios, 0.0)
        return normalized
    normalized = jnp.where(valid, jnp.log(ratios), 0.0)
    return normalized

fill_band_nans 

fill_band_nans(band: ndarray | ndarray, *, min_neighbors: int = 4) -> ndarray

Replaces NaNs in the valid upper-triangular portion of a band with the row mean.

Source code in dlem/core.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
def fill_band_nans(
    band: np.ndarray | jnp.ndarray,
    *,
    min_neighbors: int = 4,
) -> jnp.ndarray:
    "Replaces NaNs in the valid upper-triangular portion of a band with the row mean."
    if min_neighbors < 0 or min_neighbors > 8:
        raise ValueError("min_neighbors must be between 0 and 8.")

    arr = np.asarray(band, dtype=np.float32)
    width, n = arr.shape
    fallback = np.asarray(_fill_band_nan_with_row_mean(arr), dtype=np.float32)
    result = np.array(arr, copy=True)

    def _lookup(source: np.ndarray, row_idx: int, col_idx: int) -> float:
        if not (0 <= row_idx < n and 0 <= col_idx < n):
            return np.nan
        if row_idx > col_idx:
            row_idx, col_idx = col_idx, row_idx
        diag = col_idx - row_idx
        if diag >= width:
            return np.nan
        value = source[diag, col_idx]
        return float(value) if np.isfinite(value) else np.nan

    for diag in range(width):
        start_col = diag
        if start_col >= n:
            break
        for col in range(start_col, n):
            if not np.isnan(arr[diag, col]):
                continue
            row_idx = col - diag
            col_idx = col
            neighbors: list[float] = []
            for d_row in (-1, 0, 1):
                for d_col in (-1, 0, 1):
                    if d_row == 0 and d_col == 0:
                        continue
                    n_row = row_idx + d_row
                    n_col = col_idx + d_col
                    val = _lookup(arr, n_row, n_col)
                    if np.isfinite(val):
                        neighbors.append(val)
            if len(neighbors) >= min_neighbors:
                result[diag, col] = float(np.mean(neighbors))
            else:
                result[diag, col] = fallback[diag, col]

    return jnp.asarray(result, dtype=jnp.float32)

_compute_rights_table 

_compute_rights_table(p_right, mat_rows)

Function to precompute right probabilities for all rows.

Source code in dlem/core.py 
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
def _compute_rights_table(p_right, mat_rows):
    "Function to precompute right probabilities for all rows."
    n = p_right.shape[0]
    if mat_rows is None:
        limit = n
    else:
        limit = max(min(int(mat_rows), n), 0)
    length = max(limit - 1, 0)
    if length == 0:
        return jnp.zeros((0, n), dtype=p_right.dtype)

    def scan_body(this_right, _):
        next_right = jnp.concatenate((this_right[-1:], this_right[:-1]))
        return next_right, this_right

    _, rights = lax.scan(scan_body, p_right, xs=None, length=length)
    return rights

_apply_vmap_update 

_apply_vmap_update(prev, p_left, rights_table, slowdown, start_row)

Function to apply dLEM update to multiple rows via vmap.

Source code in dlem/core.py 
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
def _apply_vmap_update(prev, p_left, rights_table, slowdown, start_row):
    "Function to apply dLEM update to multiple rows via vmap."
    n_rows = prev.shape[0]
    n = p_left.shape[0]
    max_rows = max(n_rows - 1 - start_row, 0)
    max_cols = max(n - 1 - start_row, 0)
    max_rights = max(rights_table.shape[0] - start_row, 0)
    row_count = min(max_rows, max_cols, max_rights)
    if row_count == 0:
        return prev

    pl_l = p_left[:-1]
    pl_r = p_left[1:]
    rows = prev[start_row : start_row + row_count]
    rights = rights_table[start_row : start_row + row_count]

    def update(row, right):
        nom = row[:-1] * pl_l + row[1:] * right[1:]
        denom = pl_r + right[:-1] + slowdown
        return nom / denom

    new_rows = vmap(update)(rows, rights)
    return prev.at[start_row + 1 : start_row + 1 + row_count, 1:].set(new_rows)

Training utilities

_compute_rights_table 

_compute_rights_table(p_right, mat_rows)

Function to precompute right probabilities for all rows.

Source code in dlem/core.py 
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
def _compute_rights_table(p_right, mat_rows):
    "Function to precompute right probabilities for all rows."
    n = p_right.shape[0]
    if mat_rows is None:
        limit = n
    else:
        limit = max(min(int(mat_rows), n), 0)
    length = max(limit - 1, 0)
    if length == 0:
        return jnp.zeros((0, n), dtype=p_right.dtype)

    def scan_body(this_right, _):
        next_right = jnp.concatenate((this_right[-1:], this_right[:-1]))
        return next_right, this_right

    _, rights = lax.scan(scan_body, p_right, xs=None, length=length)
    return rights

_apply_vmap_update 

_apply_vmap_update(prev, p_left, rights_table, slowdown, start_row)

Function to apply dLEM update to multiple rows via vmap.

Source code in dlem/core.py 
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
def _apply_vmap_update(prev, p_left, rights_table, slowdown, start_row):
    "Function to apply dLEM update to multiple rows via vmap."
    n_rows = prev.shape[0]
    n = p_left.shape[0]
    max_rows = max(n_rows - 1 - start_row, 0)
    max_cols = max(n - 1 - start_row, 0)
    max_rights = max(rights_table.shape[0] - start_row, 0)
    row_count = min(max_rows, max_cols, max_rights)
    if row_count == 0:
        return prev

    pl_l = p_left[:-1]
    pl_r = p_left[1:]
    rows = prev[start_row : start_row + row_count]
    rights = rights_table[start_row : start_row + row_count]

    def update(row, right):
        nom = row[:-1] * pl_l + row[1:] * right[1:]
        denom = pl_r + right[:-1] + slowdown
        return nom / denom

    new_rows = vmap(update)(rows, rights)
    return prev.at[start_row + 1 : start_row + 1 + row_count, 1:].set(new_rows)

_forward_scan 

_forward_scan(p_left, p_right, mat, slowdown, start_row, steps, collect=False)

Function to perform multiple dLEM updates via lax.scan.

Source code in dlem/core.py 
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
def _forward_scan(
    p_left,
    p_right,
    mat,
    slowdown,
    start_row,
    steps,
    collect=False,
):
    "Function to perform multiple dLEM updates via lax.scan."
    start_row = core.concrete_or_error(
        int, start_row, "start_row must be a static integer."
    )
    steps = core.concrete_or_error(
        int, steps, "steps must be a static integer for _forward_scan."
    )
    slowdown = jnp.asarray(slowdown, dtype=mat.dtype)
    rights_table = _compute_rights_table(p_right, mat.shape[0])

    def body(prev, _):
        next_mat = _apply_vmap_update(prev, p_left, rights_table, slowdown, start_row)
        return next_mat, prev if collect else None

    final, history = lax.scan(body, mat, xs=None, length=steps)
    return final, history

_forward_impl 

_forward_impl(p_left, p_right, mat, slowdown, *, start_row, steps)

JIT-compiled dLEM forward pass implementation.

Source code in dlem/core.py 
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
@functools.partial(
    jax.jit, static_argnames=("start_row", "steps")
)
def _forward_impl(p_left, p_right, mat, slowdown, *, start_row, steps):
    "JIT-compiled dLEM forward pass implementation."
    final, _ = _forward_scan(
        p_left,
        p_right,
        mat,
        slowdown,
        start_row=start_row,
        steps=steps,
        collect=False,
    )
    return final

forward_custom 

forward_custom(p_left, p_right, mat, slowdown, *, start_row=0, steps=1)

Custom dLEM forward pass with adjustable start row and steps.

Source code in dlem/core.py 
270
271
272
273
274
275
276
277
278
279
def forward_custom(p_left, p_right, mat, slowdown, *, start_row=0, steps=1):
    "Custom dLEM forward pass with adjustable start row and steps."
    return _forward_impl(
        p_left,
        p_right,
        mat,
        slowdown,
        start_row=start_row,
        steps=steps,
    )

jax_forward_generate 

jax_forward_generate(p_left, p_right, slowdown, rows)

Basic dLEM forward pass that produces a full matrix of row updates (first row all ones).

Source code in dlem/core.py 
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
def jax_forward_generate(p_left, p_right, slowdown, rows):
    "Basic dLEM forward pass that produces a full matrix of row updates (first row all ones)."
    p_left = jnp.asarray(p_left, dtype=jnp.float32)
    p_right = jnp.asarray(p_right, dtype=jnp.float32)
    slowdown = jnp.asarray(slowdown, dtype=p_left.dtype)
    rows = int(rows)
    if rows <= 0:
        raise ValueError("rows must be positive")

    n = p_left.shape[0]
    initial_row = jnp.ones((n,), dtype=p_left.dtype)

    def scan_body(carry, _):
        prev_row, this_right = carry
        nom = prev_row[:-1] * p_left[:-1] + prev_row[1:] * this_right[1:]
        denom = p_left[1:] + this_right[:-1] + slowdown
        updated = nom / denom
        next_row = prev_row.at[1:].set(updated)
        next_right = jnp.concatenate((this_right[-1:], this_right[:-1]))
        return (next_row, next_right), next_row

    length = max(rows - 1, 0)
    (_, _), tail = lax.scan(scan_body, (initial_row, p_right), xs=None, length=length)
    if rows == 1:
        return initial_row[None, :]
    return jnp.vstack((initial_row, tail))