cyjax.util#

pop(arr, index)

Pop value of arr at index in JIT-compatible way.

insert_1d(arr, val, index)

Insert value into array at given index.

insert_col(mat, col, index)

Insert a column or column of sub-rows into a 2D matrix.

shuffle_axis(key, x, axis)

Shuffle indices of one axis.

binomial(n, k)

Binomial coefficients (cannot apply JIT).

flatten_coord(zs[, zs_c, need_c])

Flatten coordinates and return old shape.

unflatten_coord(arr, old_shape)

Bring output arrays to old shape before flatten_coord.

mc_integrate(key, count, fn, sample[, var])

Monte Carlo integration.

mc_integrate_batched(key, batches, ...)

Monte carlo integration in batches.