interface Module

interface Module#

This module defines a friendly interface for creating models in the Pangolin IR.

Example

Take the following model:

x    ~ normal(0,1)
y[i] ~ exponential(d[i])
z[i] ~ normal(x, y[i])

In Pangolin, you can declare this model like so:

>>> from pangolin import interface as pi
>>>
>>> x = pi.normal(0,1)
>>> y = pi.vmap(pi.exponential)(pi.constant([2,3,4]))
>>> z = pi.vmap(pi.normal, [None, 0])(x, y)

This produces the following internal representation.

>>> pi.print_upstream(z)
shape | statement
----- | ---------
()    | a = 0
()    | b = 1
()    | c ~ normal(a,b)
(3,)  | d = [2 3 4]
(3,)  | e ~ vmap(exponential, [0], 3)(d)
(3,)  | f ~ vmap(normal, [None, 0], 3)(c,e)

Reference card#

The entire interface is based on methods that create InfixRV objects (or that create functions that create InfixRV objects).

Here are all the functions:

To do…	Use…
Constants	`constant()`
Arithmetic	`add()` `sub()` `mul()` `div()` (or infix `+`, `-`, `*`, `/`)
Trigonometry	`arccos()` `arccosh()` `arcsin()` `arcsinh()` `arctan()` `arctanh()` `cos()` `cosh()` `sin()` `sinh()` `tan()` `tanh()`
Other scalar functions	`pow()` (or infix `**`) `sqrt()` `abs()` `exp()` `inv_logit()` `expit()` `sigmoid()` `log()` `loggamma()` `logit()` `step()`
Linear algebra	`matmul()` (or infix `@`) `transpose` (or `rv.T`) `inv` `cholesky`
Other multivariate functions	`sum()` `softmax()` `diag` `diag_matrix` `fill_tril` `extract_tril`
Scalar distributions	`normal()` `normal_prec()` `lognormal()` `cauchy()` `bernoulli()` `bernoulli_logit()` `beta()` `binomial()` `categorical()` `uniform()` `beta_binomial()` `exponential()` `gamma()` `poisson()` `student_t()`
Multivariate distributions	`multi_normal()` `multinomial()` `dirichlet()` `wishart()`
Control flow	`vmap()` `broadcast` `composite()` `autoregressive()` `autoregress()`
Indexing	`index()` (or `rv.[...]`) `vindex`
Transformed densities	`Transform` `tforms` `compose_transforms`

Auto-casting#

Most functions take RVLike arguments, meaning either an RV or something that can be cast to a NumPy array. In the latter case, it is implicitly cast to a constant RV. For example, categorical takes an RVLike argument, so instead of tediously writin this:

>>> probs_list = [0.1, 0.2, 0.7]
>>> probs_array = np.array(probs_list)
>>> probs_rv = constant(probs_array)
>>> x = categorical(probs_rv)

You can simply write:

>>> x = categorical([0.1, 0.2, 0.7])

Broadcasting#

Do you love broadcasting? Do you hate broadcasting? Are you lukewarm about broadcasting? Well, good news. In this interface, you can configure how broadcasting works. You can use “simple” broadcasting, you can use full NumPy-style broadcasting, or you can turn broadcasting off completely. See Broadcasting for details.

API docs#

pangolin.interface.constant(value)[source]#

Create a constant RV

Parameters:: value (ArrayLike) – value for the constant. Should be a numpy (or JAX) array or something castable to that, e.g. int / float / list of list of ints/floats.
Returns:: out – RV with Constant Op
Return type:: InfixRV[Constant]

Examples

>>> constant(7)
InfixRV(Constant(7))
>>> constant([0,1,2])
InfixRV(Constant([0,1,2]))

pangolin.interface.vmap(f, in_axes=0, axis_size=None)[source]#

Vectorizing map. Create a function which maps f over argument axes.

This function matches exactly the interface of jax.vmap, although it doesn’t provide some of the optional arguments jax.vmap does.

Parameters:

f (Callable[[*Args], PyTree[InfixRV]]) – The function to vmap. Should take a pytree of RV as inputs and return a pytree of RV as outputs.
in_axes (Any) – An int, None, or pytree with roots that are int or None. Specifies which axis of each RV should be mapped (if int) or that no axis shuld be mapped (if None). Can be a pytree matching the structure of all arguments to f. Or, can be a pytree that is a prefix to the pytree representing all arguments. By default, in_axes is zero, meaning all RVs are mapped over the first axis.
axis_size (int | None) – An integer indicating the size of the axis to be mapped. This is optional unless all leaves of in_axes are None.

Returns:

vec_f – batched/vectorized version of f with arguments matching those of f with extra axes at positions indicated by in_axes and a return value that corresponds to that of f but with an extra axis in the first position.

Return type:

Callable[[*Args], PyTree[InfixRV]]

Examples

Here’s the simplest possible example.

>>> def fun(a):
...     return exp(a)
>>> A = constant([0,1,2])
>>> vmap(fun)(A)
InfixRV(VMap(Exp(), [0], 3), InfixRV(Constant([0,1,2])))

Multiple inputs are OK.

>>> def fun(a,b):
...     return a*b
>>> A = constant([1,2,3])
>>> B = constant([4,5,6])
>>> vmap(fun)(A, B)
InfixRV(VMap(Mul(), [0, 0], 3), InfixRV(Constant([1,2,3])), InfixRV(Constant([4,5,6])))

Unmapped inputs are OK.

>>> def fun(a,b):
...     return a*b
>>> A = constant([1,2,3])
>>> vmap(fun, [0, None])(A, constant(7))
InfixRV(VMap(Mul(), [0, None], 3), InfixRV(Constant([1,2,3])), InfixRV(Constant(7)))

Multiple outputs are OK.

>>> def fun(a):
...     return [exp(a), log(a)]
>>> [out1, out2] = vmap(fun)(A)
>>> out1
InfixRV(VMap(Exp(), [0], 3), InfixRV(Constant([1,2,3])))
>>> out2
InfixRV(VMap(Log(), [0], 3), InfixRV(Constant([1,2,3])))

Pytree inputs and pytree in_axes are OK

>>> def fun(x):
...     return x['cat']*x['dog']
>>> x = {'cat': A, 'dog': constant(3)}
>>> in_axes = {'cat': 0, 'dog': None}
>>> vmap(fun, in_axes)(x)
InfixRV(VMap(Mul(), [0, None], 3), InfixRV(Constant([1,2,3])), InfixRV(Constant(3)))

Pytree outputs are OK

>>> def fun(a, b):
...     return {"add": a+b, "mul": a*b}
>>> vmap(fun)(A, B)
{'add': InfixRV(VMap(Add(), [0, 0], 3), InfixRV(Constant([1,2,3])), InfixRV(Constant([4,5,6]))), 'mul': InfixRV(VMap(Mul(), [0, 0], 3), InfixRV(Constant([1,2,3])), InfixRV(Constant([4,5,6])))}

Pytree in_axis prefixes are OK

>>> def fun(x):
...     [a, (b,c)] = x
...     return (a*b)+c
>>> x = [A, (constant(7), constant(8))]
>>> in_axes1 = [0, (None, None)] # axes for each leaf
>>> in_axes2 = [0, None]         # single None for (b,c) tuple!
>>> out1 = vmap(fun, in_axes1)(x)
>>> out2 = vmap(fun, in_axes2)(x)
>>> out1.op == out2.op
True
>>> print_upstream(out1)
shape | statement
----- | ---------
(3,)  | a = [1 2 3]
()    | b = 7
(3,)  | c = vmap(mul, [0, None], 3)(a,b)
()    | d = 8
(3,)  | e = vmap(add, [0, None], 3)(c,d)
>>> print_upstream(out2)
shape | statement
----- | ---------
(3,)  | a = [1 2 3]
()    | b = 7
(3,)  | c = vmap(mul, [0, None], 3)(a,b)
()    | d = 8
(3,)  | e = vmap(add, [0, None], 3)(c,d)

See also

pangolin.ir.VMap

pangolin.interface.broadcast(scalar_fun)[source]#

Parameters:: scalar_fun – Function accepting and returning all-scalar arguments. Can have keyword arguments, but cannot have keyword-only arguments.
Returns:: “broadcast” version of scalar_fun that accepts and returns array-shaped random variables according to the current broadcasting rules.

Examples

Take a simple function with scalar input and output

>>> fun = lambda a: InfixRV(ir.Exponential(), a)

If you just have a scalar input, this function works fine in all three broadcasting modes

>>> with override(broadcasting='off'):
...     broadcast(fun)(3)
InfixRV(Exponential(), InfixRV(Constant(3)))
>>> with override(broadcasting='simple'):
...      broadcast(fun)(3)
InfixRV(Exponential(), InfixRV(Constant(3)))
>>> with override(broadcasting='numpy'):
...      broadcast(fun)(3)
InfixRV(Exponential(), InfixRV(Constant(3)))

If you have a vector input, this only works with simple or numpy broadcasting

>>> with override(broadcasting='off'):
...     broadcast(fun)([3,4])
Traceback (most recent call last):
...
ValueError: Non-scalar argument: [3 4]
>>> with override(broadcasting='simple'):
...     broadcast(fun)([3,4])
InfixRV(VMap(Exponential(), [0], 2), InfixRV(Constant([3,4])))
>>> with override(broadcasting='numpy'):
...     broadcast(fun)([3,4])
InfixRV(VMap(Exponential(), [0], 2), InfixRV(Constant([3,4])))

Now take this function with one vector and one matrix argument.

>>> fun = lambda a, b: InfixRV(ir.Normal(), a, b)
>>> mean = [1,2,3]
>>> scale = [[1,2,3],[4,5,6]]

This will only work with numpy-style broadcasting

>>> with override(broadcasting='off'):
...     broadcast(fun)(mean, scale)
Traceback (most recent call last):
...
ValueError: Non-scalar argument: [1 2 3]
>>> with override(broadcasting='simple'):
...     broadcast(fun)(mean, scale)
Traceback (most recent call last):
...
ValueError: Can't broadcast non-matching shapes (2, 3) and (3,)
>>> with override(broadcasting='numpy'):
...     z = broadcast(fun)(mean, scale)
>>> print(z)
vmap(vmap(normal, [0, 0], 3), [None, 0], 2)([1 2 3], [[1 2 3] [4 5 6]])

pangolin.interface.composite(fun)[source]#

Turn a function into a new function that returns a single Composite RV. Typically this would be used to create autoregressive distributions via the autoregressive function, rather than called directly by the user.

Parameters:: fun (Callable) – Function that takes any number of RVs or pytrees of RVs and returns a single RV
Return type:: Callable[[…], InfixRV[Composite]]

Examples

>>> def f(stuff):
...     x = stuff['cat']
...     y = stuff['dog']
...     a = x+y
...     return normal(a, y)
>>> composite(f)
<function composite.<locals>.myfun at ...>
>>> x = makerv(1)
>>> y = makerv(2)
>>> z = composite(f)({'cat':x, 'dog': y})
>>> z.op
Composite(2, (Add(), Normal()), [[0, 1], [2, 1]])
>>> z.parents == (x, y)
True

>>> # Example where x is not explicitly passed
>>> def f(stuff):
...     y = stuff['dog']
...     a = x+y
...     return normal(a, y)
>>> z = composite(f)({'dog':y})
>>> z.op
Composite(2, (Add(), Normal()), [[1, 0], [2, 0]])
>>> z.parents == (y, x)
True

type pangolin.interface.Autoregressable = Callable[[InfixRV], InfixRV] | Callable[[InfixRV, PyTree[InfixRV]], InfixRV] | Callable[[InfixRV, PyTree[InfixRV], PyTree[InfixRV]], InfixRV] | Callable[[InfixRV, PyTree[InfixRV], PyTree[InfixRV], PyTree[InfixRV]], InfixRV] | Callable[[InfixRV, PyTree[InfixRV], PyTree[InfixRV], PyTree[InfixRV], PyTree[InfixRV]], InfixRV] | Callable[[InfixRV, PyTree[InfixRV], PyTree[InfixRV], PyTree[InfixRV], PyTree[InfixRV], PyTree[InfixRV]], InfixRV] | Callable[[InfixRV, PyTree[InfixRV], PyTree[InfixRV], PyTree[InfixRV], PyTree[InfixRV], PyTree[InfixRV], PyTree[InfixRV]], InfixRV]#

type pangolin.interface.Autoregressed = Callable[[...], InfixRV[Autoregressive[Composite]]]#

pangolin.interface.autoregressive(fun, length=None, in_axes=0)[source]#

Given a function, create a function to generate an RV with an Autoregressive Op. Doing

auto_fun = autoregressive(fun, 5)
z = auto_fun(start)

is semantically like

carry = start
values = []
for i in range(5):
    carry = fun(carry)
    values.append(carry)
z = concatenate(values)

As a more complex example, doing

auto_fun = autoregressive(fun, 5, [0, None])
z = auto_fun(start, A, B)

is semantically like:

carry = start
values = []
# length = 5
for i in range(5):
    carry = fun(carry, A[i], B) # A sliced on 0th axis, B unsliced
    values.append(carry)
z = concatenate(values)

Even more generally, in_axes can be any pytree of in-axes that is a tree prefix for the arguments. (See examples below)

Parameters:

fun (Autoregressable[O]) – Function to call repeatedly to define the distribution. Must take some argument carry as the first input and return a single RV with the same shape as carry. This function can only create a single random RV which (if it exists) must be the return value. But it can create an arbitrary number of non-random RV internally. This function can take extra inputs, which can be RVLike or pytrees of RVLike. These can be mapped as determined Can optionally take extra inputs that will be mapped.
length (None | int) – Length of autoregressive. Can be None if any inputs are mapped along some axis.
in_axes (PyTree[int | None]) – What axis to map each input other than carry over (or None if non-mapped). As with vmap, can be a pytree of RV corresponding to the structure of all inputs other than carry.

Returns:

Function that takes a single init RVLike plus some number of pytrees of RVLike with mapped axes and produces a single RV[Autoregressive]

Return type:

Autoregressed[O]

Examples

Distribution where z[i] ~ normal(exp(z[i-1]), 1).

>>> x = constant(3.3)
>>> def fun(carry):
...     return normal(exp(carry), 1)
>>> z = autoregressive(fun, 5)(x)
>>> isinstance(z.op, Autoregressive)
True
>>> z.op.base_op
Composite(1, (Exp(), Constant(1), Normal()), [[0], [], [1, 2]])
>>> z.op.length
5
>>> z.op.in_axes
()
>>> z.parents == (x,)
True

Distribution where z[i] ~ exponential(z[i-1]*y[i])

>>> x = constant(3.3)
>>> y = constant([1,2,3])
>>> def fun(carry, yi):
...     return exponential(carry * yi)
>>> z = autoregressive(fun)(x,y)
>>> isinstance(z.op, Autoregressive)
True
>>> z.op.base_op
Composite(2, (Mul(), Exponential()), [[0, 1], [2]])
>>> z.op.length # note this was inferred!
3
>>> z.op.in_axes
(0,)
>>> z.parents == (x, y)
True

You can also pass bare constants

>>> def fun(carry, yi):
...     return exponential(carry * yi)
>>> z = autoregressive(fun)(3.3, np.array([1,2,3]))
>>> isinstance(z.op, Autoregressive)
True
>>> z.op.base_op
Composite(2, (Mul(), Exponential()), [[0, 1], [2]])
>>> z.op.length
3
>>> z.op.in_axes
(0,)

pangolin.interface.autoregress(length=None, in_axes=0)[source]#

Simple decorator to create functions to create autoregressive RVs. The idea is that

autoregress(length, in_axes)(fun)

is exactly the same as

autoregressive(fun, length, in_axes)

This can be very convenient as a decorator.

Parameters:

length (int | None) – the number of repetitions
in_axes (Any) – axis to map arguments other than carry over

Returns:

Decorator that takes a function that transforms a RVLike and some number of pytress of RVLike into a single RV and returns a function that transforms a RVLike and some number of pytrees of RVLike into a single RV with an ir.Autoregressive op.

Return type:

Callable[[Autoregressable], Autoregressed]

Examples

Here are two equivalent examples. Here’s autoregressive:

>>> x = constant(3.3)
>>> def fun(carry):
...     return normal(exp(carry), 1)
>>> z1 = autoregressive(fun,5)(x)

And here’s autoregress:

>>> @autoregress(5)
... def fun(carry):
...     return normal(exp(carry), 1)
>>> z2 = fun(x)

>>> z1.op == z2.op
True
>>> z1.parents == z2.parents == (x,)
True

pangolin.interface.index(var, *indices)[source]#

Index a RV Using fully-orthogonal indexing.

Note that this function is (intentionally) much simpler than indexing in NumPy or JAX or PyTorch in that it performs fully orthogonal indexing and that slices are treated exactly the same as 1D arrays.

(Similar to oindex in NEP 21 .)

Parameters:

var (RVLike) – The RV to be indexed
indices (ArrayLike | InfixRV | slice | EllipsisType) – The indices into the RV. Unless there is an ellipsis, the number of indices must match the number of dimensions of var.

Returns:

Random variable with shape equal to the shapes of all indices, concatenated.

Return type:

InfixRV[Index]

Examples

>>> A = constant([[3,0,2],[4,4,4]])
>>> B = index(A, slice(None), [2,2])
>>> print_upstream(B)
shape  | statement
------ | ---------
(2, 3) | a = [[3 0 2] [4 4 4]]
(2,)   | b = [0 1]
(2,)   | c = [2 2]
(2, 2) | d = index(a,b,c)
>>> C = index(A, 0, ...)
>>> print_upstream(C)
shape  | statement
------ | ---------
(2, 3) | a = [[3 0 2] [4 4 4]]
()     | b = 0
(3,)   | c = [0 1 2]
(3,)   | d = index(a,b,c)

Technically, it’s legal (although pointless) to index a 0-D array

>>> A = constant(12.0)
>>> B = index(A)
>>> print_upstream(B)
shape | statement
----- | ---------
()    | a = 12.
()    | b = index(a)

pangolin.interface.vindex(var, *indices)[source]#

Index a RV Using “fully-vectorized” indexing.

This function treats var as a scalar function of the indices and then applys the normal scalar broadcasting rules (as configured by config). Note that this behavior is (intentionally) much simpler than indexing in NumPy or JAX or PyTorch in that it performs fully orthogonal indexing and that slices are treated exactly the same as 1D arrays.

(Named in honor of vindex in NEP 21 .)

Parameters:

var (RVLike) – The RV to be indexed
indices (ArrayLike | InfixRV | slice | EllipsisType) – The indices into the RV. Unless there is an ellipsis, the number of indices must match the number of dimensions of var.

Returns:

Random variable with shape equal to the shapes of all indices, concatenated.

Return type:

InfixRV[Index] | InfixRV[VMap[Index]]

Examples

Index a 2D array with two matching 2D arrays

>>> A = constant([[0.,0,0],[1,1,1],[2,2,2]])
>>> B = constant([[0,1,2],[2,1,0]])
>>> C = constant([[2,1,1],[0,0,0]])
>>> D = vindex(A, B, C)
>>> print_upstream(D)
shape  | statement
------ | ---------
(3, 3) | a = [[0. 0. 0.] [1. 1. 1.] [2. 2. 2.]]
(2, 3) | b = [[0 1 2] [2 1 0]]
(2, 3) | c = [[2 1 1] [0 0 0]]
(2, 3) | d = vmap(vmap(index, [None, 0, 0], 3), [None, 0, 0], 2)(a,b,c)

With a square matrix, you can use slices to extract the diagonal

>>> A = constant([[1,2,3],[4,5,6],[7,8,9]])
>>> B = vindex(A,slice(None),slice(None))
>>> print_upstream(B)
shape  | statement
------ | ---------
(3, 3) | a = [[1 2 3] [4 5 6] [7 8 9]]
(3,)   | b = [0 1 2]
(3,)   | c = [0 1 2]
(3,)   | d = vmap(index, [None, 0, 0], 3)(a,b,c)

You can also get the anti-diagonal

>>> C = vindex(A,slice(None),slice(None,None,-1))
>>> print_upstream(C)
shape  | statement
------ | ---------
(3, 3) | a = [[1 2 3] [4 5 6] [7 8 9]]
(3,)   | b = [0 1 2]
(3,)   | c = [2 1 0]
(3,)   | d = vmap(index, [None, 0, 0], 3)(a,b,c)

pangolin.interface.normal(mu, sigma)[source]#

Creates a Normal distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

mu (RVLike) – location / mean
sigma (RVLike) – scale / standard deviation

Returns:

Random variable with z.op of type pangolin.ir.Normal or pangolin.ir.VMap if broadcasting is triggered and 2 parent(s).

Return type:

InfixRV[Normal | VMap[Normal]]

Examples

>>> normal(0.1, 0.2)
InfixRV(Normal(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)))

pangolin.interface.normal_prec(mu, tau)[source]#

Creates a NormalPrec distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

mu (RVLike) – location / mean
tau (RVLike) – precision / inverse variance

Returns:

Random variable with z.op of type pangolin.ir.NormalPrec or pangolin.ir.VMap if broadcasting is triggered and 2 parent(s).

Return type:

InfixRV[NormalPrec | VMap[NormalPrec]]

Examples

>>> normal_prec(0.1, 0.2)
InfixRV(NormalPrec(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)))

pangolin.interface.lognormal(mu, sigma)[source]#

Creates a Lognormal distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

mu (RVLike) – logarithm of location
sigma (RVLike) – logarithm of scale (not sigma squared!)

Returns:

Random variable with z.op of type pangolin.ir.Lognormal or pangolin.ir.VMap if broadcasting is triggered and 2 parent(s).

Return type:

InfixRV[Lognormal | VMap[Lognormal]]

Examples

>>> lognormal(0.1, 0.2)
InfixRV(Lognormal(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)))

pangolin.interface.cauchy(mu, sigma)[source]#

Creates a Cauchy distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

mu (RVLike) – location
sigma (RVLike) – scale

Returns:

Random variable with z.op of type pangolin.ir.Cauchy or pangolin.ir.VMap if broadcasting is triggered and 2 parent(s).

Return type:

InfixRV[Cauchy | VMap[Cauchy]]

Examples

>>> cauchy(0.1, 0.2)
InfixRV(Cauchy(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)))

pangolin.interface.bernoulli(theta)[source]#

Creates a Bernoulli distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: theta (RVLike) – probability (between 0 and 1)
Returns:: Random variable with z.op of type pangolin.ir.Bernoulli or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Bernoulli | VMap[Bernoulli]]

Examples

>>> bernoulli(0.1)
InfixRV(Bernoulli(), InfixRV(Constant(0.1)))

pangolin.interface.bernoulli_logit(theta)[source]#

Creates a BernoulliLogit distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: theta (RVLike) – logit of probability (unbounded)
Returns:: Random variable with z.op of type pangolin.ir.BernoulliLogit or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[BernoulliLogit | VMap[BernoulliLogit]]

Examples

>>> bernoulli_logit(0.1)
InfixRV(BernoulliLogit(), InfixRV(Constant(0.1)))

pangolin.interface.binomial(N, theta)[source]#

Creates a Binomial distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

N (RVLike) – number of trials
theta (RVLike) – probability of success for each trial

Returns:

Random variable with z.op of type pangolin.ir.Binomial or pangolin.ir.VMap if broadcasting is triggered and 2 parent(s).

Return type:

InfixRV[Binomial | VMap[Binomial]]

Examples

>>> binomial(0.1, 0.2)
InfixRV(Binomial(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)))

pangolin.interface.uniform(alpha, beta)[source]#

Creates a Uniform distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

alpha (RVLike) – lower bound
beta (RVLike) – upper bound

Returns:

Random variable with z.op of type pangolin.ir.Uniform or pangolin.ir.VMap if broadcasting is triggered and 2 parent(s).

Return type:

InfixRV[Uniform | VMap[Uniform]]

Examples

>>> uniform(0.1, 0.2)
InfixRV(Uniform(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)))

pangolin.interface.beta(alpha, beta)[source]#

Creates a Beta distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

alpha (RVLike) – shape
beta (RVLike) – shape

Returns:

Random variable with z.op of type pangolin.ir.Beta or pangolin.ir.VMap if broadcasting is triggered and 2 parent(s).

Return type:

InfixRV[Beta | VMap[Beta]]

Examples

>>> beta(0.1, 0.2)
InfixRV(Beta(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)))

pangolin.interface.beta_binomial(N, alpha, beta)[source]#

Creates a BetaBinomial distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

N (RVLike) – as in binomial dist
alpha (RVLike) – as in beta dist
beta (RVLike) – as in beta dist

Returns:

Random variable with z.op of type pangolin.ir.BetaBinomial or pangolin.ir.VMap if broadcasting is triggered and 3 parent(s).

Return type:

InfixRV[BetaBinomial | VMap[BetaBinomial]]

Examples

>>> beta_binomial(0.1, 0.2, 0.4)
InfixRV(BetaBinomial(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)), InfixRV(Constant(0.4)))

Notes

This follows the (N,alpha,beta) convention of Stan (and Wikipedia). Some other systems (e.g. Numpyro) use alternate variable orderings. This is no problem for you as a user, since pangolin does the re-ordering for you based on the backend. But keep it in mind if translating a model from one system to another.

pangolin.interface.exponential(beta)[source]#

Creates a Exponential distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: beta (RVLike) – rate / inverse scale
Returns:: Random variable with z.op of type pangolin.ir.Exponential or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Exponential | VMap[Exponential]]

Examples

>>> exponential(0.1)
InfixRV(Exponential(), InfixRV(Constant(0.1)))

pangolin.interface.gamma(alpha, beta)[source]#

Creates a Gamma distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

alpha (RVLike) – shape
beta (RVLike) – rate / inverse scale

Returns:

Random variable with z.op of type pangolin.ir.Gamma or pangolin.ir.VMap if broadcasting is triggered and 2 parent(s).

Return type:

InfixRV[Gamma | VMap[Gamma]]

Examples

>>> gamma(0.1, 0.2)
InfixRV(Gamma(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)))

Notes

This follows Stan in using the “shape/rate” parameterization, not the “shape/scale” parameterization.

pangolin.interface.poisson(lambd)[source]#

Creates a Poisson distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: lambd (RVLike) – lambda
Returns:: Random variable with z.op of type pangolin.ir.Poisson or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Poisson | VMap[Poisson]]

Examples

>>> poisson(0.1)
InfixRV(Poisson(), InfixRV(Constant(0.1)))

pangolin.interface.student_t(nu, mu, sigma)[source]#

Creates a StudentT distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

nu (RVLike) – degress of freedom
mu (RVLike) – location (often 0)
sigma (RVLike) – scale (often 1)

Returns:

Random variable with z.op of type pangolin.ir.StudentT or pangolin.ir.VMap if broadcasting is triggered and 3 parent(s).

Return type:

InfixRV[StudentT | VMap[StudentT]]

Examples

>>> student_t(0.1, 0.2, 0.4)
InfixRV(StudentT(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)), InfixRV(Constant(0.4)))

Also, let’s check that we can mix positional and keyword arguments.

>>> student_t(nu=0.1, mu=0.2, sigma=0.4)
InfixRV(StudentT(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)), InfixRV(Constant(0.4)))

>>> student_t(0.1, mu=0.2, sigma=0.4)
InfixRV(StudentT(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)), InfixRV(Constant(0.4)))

>>> student_t(0.1, sigma=0.4, mu=0.2)
InfixRV(StudentT(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)), InfixRV(Constant(0.4)))

>>> student_t(sigma=0.4, nu=0.1, mu=0.2)
InfixRV(StudentT(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)), InfixRV(Constant(0.4)))

pangolin.interface.add(a, b)[source]#

Creates a Add distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

a (RVLike) – parent 0
b (RVLike) – parent 1

Returns:

Random variable with z.op of type pangolin.ir.Add or pangolin.ir.VMap if broadcasting is triggered and 2 parent(s).

Return type:

InfixRV[Add | VMap[Add]]

Examples

>>> add(0.1, 0.2)
InfixRV(Add(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)))

pangolin.interface.sub(a, b)[source]#

Creates a Sub distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

a (RVLike) – parent 0
b (RVLike) – parent 1

Returns:

Random variable with z.op of type pangolin.ir.Sub or pangolin.ir.VMap if broadcasting is triggered and 2 parent(s).

Return type:

InfixRV[Sub | VMap[Sub]]

Examples

>>> sub(0.1, 0.2)
InfixRV(Sub(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)))

pangolin.interface.mul(a, b)[source]#

Creates a Mul distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

a (RVLike) – parent 0
b (RVLike) – parent 1

Returns:

Random variable with z.op of type pangolin.ir.Mul or pangolin.ir.VMap if broadcasting is triggered and 2 parent(s).

Return type:

InfixRV[Mul | VMap[Mul]]

Examples

>>> mul(0.1, 0.2)
InfixRV(Mul(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)))

pangolin.interface.div(a, b)[source]#

Creates a Div distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

a (RVLike) – parent 0
b (RVLike) – parent 1

Returns:

Random variable with z.op of type pangolin.ir.Div or pangolin.ir.VMap if broadcasting is triggered and 2 parent(s).

Return type:

InfixRV[Div | VMap[Div]]

Examples

>>> div(0.1, 0.2)
InfixRV(Div(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)))

pangolin.interface.pow(a, b)[source]#

Creates a Pow distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:

a (RVLike) – parent 0
b (RVLike) – parent 1

Returns:

Random variable with z.op of type pangolin.ir.Pow or pangolin.ir.VMap if broadcasting is triggered and 2 parent(s).

Return type:

InfixRV[Pow | VMap[Pow]]

Examples

>>> pow(0.1, 0.2)
InfixRV(Pow(), InfixRV(Constant(0.1)), InfixRV(Constant(0.2)))

pangolin.interface.sqrt(x)[source]#

Take a square root.

sqrt(x) is an alias for pow(x, 0.5)

Parameters:: x (RVLike)
Return type:: InfixRV[Pow | VMap[Pow]]

pangolin.interface.abs(a)[source]#

Creates a Abs distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Abs or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Abs | VMap[Abs]]

Examples

>>> abs(0.1)
InfixRV(Abs(), InfixRV(Constant(0.1)))

pangolin.interface.arccos(a)[source]#

Creates a Arccos distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Arccos or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Arccos | VMap[Arccos]]

Examples

>>> arccos(0.1)
InfixRV(Arccos(), InfixRV(Constant(0.1)))

pangolin.interface.arccosh(a)[source]#

Creates a Arccosh distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Arccosh or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Arccosh | VMap[Arccosh]]

Examples

>>> arccosh(0.1)
InfixRV(Arccosh(), InfixRV(Constant(0.1)))

pangolin.interface.arcsin(a)[source]#

Creates a Arcsin distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Arcsin or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Arcsin | VMap[Arcsin]]

Examples

>>> arcsin(0.1)
InfixRV(Arcsin(), InfixRV(Constant(0.1)))

pangolin.interface.arcsinh(a)[source]#

Creates a Arcsinh distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Arcsinh or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Arcsinh | VMap[Arcsinh]]

Examples

>>> arcsinh(0.1)
InfixRV(Arcsinh(), InfixRV(Constant(0.1)))

pangolin.interface.arctan(a)[source]#

Creates a Arctan distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Arctan or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Arctan | VMap[Arctan]]

Examples

>>> arctan(0.1)
InfixRV(Arctan(), InfixRV(Constant(0.1)))

pangolin.interface.arctanh(a)[source]#

Creates a Arctanh distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Arctanh or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Arctanh | VMap[Arctanh]]

Examples

>>> arctanh(0.1)
InfixRV(Arctanh(), InfixRV(Constant(0.1)))

pangolin.interface.cos(a)[source]#

Creates a Cos distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Cos or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Cos | VMap[Cos]]

Examples

>>> cos(0.1)
InfixRV(Cos(), InfixRV(Constant(0.1)))

pangolin.interface.cosh(a)[source]#

Creates a Cosh distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Cosh or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Cosh | VMap[Cosh]]

Examples

>>> cosh(0.1)
InfixRV(Cosh(), InfixRV(Constant(0.1)))

pangolin.interface.exp(a)[source]#

Creates a Exp distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Exp or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Exp | VMap[Exp]]

Examples

>>> exp(0.1)
InfixRV(Exp(), InfixRV(Constant(0.1)))

pangolin.interface.inv_logit(a)[source]#

Creates a InvLogit distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.InvLogit or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[InvLogit | VMap[InvLogit]]

Examples

>>> inv_logit(0.1)
InfixRV(InvLogit(), InfixRV(Constant(0.1)))

pangolin.interface.expit(a)[source]#: An alias for inv_logit

pangolin.interface.sigmoid(a)[source]#: An alias for inv_logit

pangolin.interface.log(a)[source]#

Creates a Log distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Log or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Log | VMap[Log]]

Examples

>>> log(0.1)
InfixRV(Log(), InfixRV(Constant(0.1)))

pangolin.interface.loggamma(a)[source]#

Creates a Loggamma distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Loggamma or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Loggamma | VMap[Loggamma]]

Examples

>>> loggamma(0.1)
InfixRV(Loggamma(), InfixRV(Constant(0.1)))

Notes

Do we want scipy.special.loggamma or scipy.special.gammaln?

pangolin.interface.logit(a)[source]#

Creates a Logit distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Logit or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Logit | VMap[Logit]]

Examples

>>> logit(0.1)
InfixRV(Logit(), InfixRV(Constant(0.1)))

pangolin.interface.sin(a)[source]#

Creates a Sin distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Sin or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Sin | VMap[Sin]]

Examples

>>> sin(0.1)
InfixRV(Sin(), InfixRV(Constant(0.1)))

pangolin.interface.sinh(a)[source]#

Creates a Sinh distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Sinh or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Sinh | VMap[Sinh]]

Examples

>>> sinh(0.1)
InfixRV(Sinh(), InfixRV(Constant(0.1)))

pangolin.interface.step(a)[source]#

Creates a Step distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Step or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Step | VMap[Step]]

Examples

>>> step(0.1)
InfixRV(Step(), InfixRV(Constant(0.1)))

pangolin.interface.tan(a)[source]#

Creates a Tan distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Tan or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Tan | VMap[Tan]]

Examples

>>> tan(0.1)
InfixRV(Tan(), InfixRV(Constant(0.1)))

pangolin.interface.tanh(a)[source]#

Creates a Tanh distributed RV.

All arguments must either be scalar or mutually broadcastable according to config.broadcasting.

Parameters:: a (RVLike) – parent 0
Returns:: Random variable with z.op of type pangolin.ir.Tanh or pangolin.ir.VMap if broadcasting is triggered and 1 parent(s).
Return type:: InfixRV[Tanh | VMap[Tanh]]

Examples

>>> tanh(0.1)
InfixRV(Tanh(), InfixRV(Constant(0.1)))

pangolin.interface.multi_normal(mean, cov)[source]#

Create a multivariate normal distributed random variable. Call as multi_normal(mean, cov)

Parameters:

mean (RVLike) – mean (1D array)
cov (RVLike) – covariance (2D positive-definite array)

Return type:

InfixRV[MultiNormal]

pangolin.interface.categorical(theta)[source]#

Create a categorical distributed RV where theta is a vector of non-negative reals that sums to one.

Parameters:: theta (RVLike) – positive event probabilities (should sum to one)
Return type:: InfixRV[Categorical]

pangolin.interface.multinomial(n, p)[source]#

Create a multinomial distributed random variable. Call as multinomial(n,p) where n is the number of repetitions and p is a vector of probabilities that sums to one.

Parameters:

n (RVLike) – number of repetitions (scalar)
p (RVLike) – vector of probabilities (should sum to one)

Return type:

InfixRV[Multinomial]

pangolin.interface.dirichlet(alpha)[source]#

Create a Dirichlet distributed random variable. Call as dirichlet(alpha) where alpha is a 1-D vector of positive reals.

Parameters:: alpha (RVLike) – concentration (vector of positive numbers)
Return type:: InfixRV[Dirichlet]

pangolin.interface.wishart(nu, S)[source]#

Create a Wishart distributed random variable.

Parameters:

nu (RVLike) – degress of freedom (scalar)
S (RVLike) – scale matrix (symmetric posisitive definite)

Return type:

InfixRV[Wishart]

pangolin.interface.matmul(a, b)[source]#

Matrix product of two arrays. The behavior follows that of numpy.matmul except that a and b must both be 1-D or 2-D arrays. In particular:

If a and b are both 1-D then this represents an inner-product.
If a is 1-D and b is 2-D then this represents vector/matrix multiplication
If a is 2-D and b is 1-D then this represents matrix/vector multiplication
If a and b are both 2-D then this represents matrix/matrix multiplication

Parameters:

a (RVLike) – first argument (1d or 2d array)
b (RVLike) – second argument (1d or 2d array, matching shape of a)

Return type:

InfixRV[Matmul]

pangolin.interface.inv(a)[source]#

Take the inverse of a matrix. Input must be a 2-D square (invertible) array.

Parameters:: a (RVLike)
Return type:: InfixRV[Inv]

pangolin.interface.transpose(a)[source]#

Tranpose a matrix. Input must be a 2-D array.

Parameters:: a (RVLike)
Return type:: InfixRV[Transpose]

pangolin.interface.cholesky(a)[source]#

Take the inverse of a matrix. Input must be a 2-D square (invertible) array.

Parameters:: a (RVLike)
Return type:: InfixRV[Cholesky]

pangolin.interface.softmax(a)[source]#

Take softmax function. (TODO: conform to syntax of scipy.special.softmax

Parameters:: a (RVLike) – 1-D vector
Return type:: InfixRV[Softmax]

pangolin.interface.sum(x, axis)[source]#

Take the sum of a random variable along a given axis

Parameters:

x (InfixRV) – an RV (or something that can be cast to a Constant RV)
axis (int) – a non-negative integer (cannot be a random variable)

Return type:

InfixRV[Sum]

pangolin.interface.diag(a)[source]#

Get the diagonal of a matrix. Input must be a 2-D square array. Does not construct diagonal matrices. (Use diag_matrix)

Parameters:: a (RVLike)
Return type:: InfixRV[Diag]

pangolin.interface.diag_matrix(a)[source]#

Get the diagonal of a matrix. Input must be a 1-D array. Does not extract diagonals. (Use diag)

Parameters:: a (RVLike)
Return type:: InfixRV[DiagMatrix]

pangolin.interface.fill_tril(params)[source]#

Take a 1D array, fill the lower-triangular entries of a matrix using Pangolin Ops. This does not create a special Op but rather reduces to indexing and multiplication.

Parameters:: params (ArrayLike | InfixRV) – 1D array of length N*(N+1)/2
Returns:: 2D Lower-triangular matrix with shape (N,N)

Examples

>>> params = constant([1.1,2.2,3.3])
>>> out = fill_tril(params)
>>> print_upstream(out)
shape  | statement
------ | ---------
(3,)   | a = [1.1 2.2 3.3]
(2, 2) | b = [[0 0] [1 2]]
(2, 2) | c = index(a,b)
(2, 2) | d = [[1. 0.] [1. 1.]]
(2, 2) | e = vmap(vmap(mul, [0, 0], 2), [0, 0], 2)(c,d)

pangolin.interface.extract_tril(L)[source]#

Given a square 2D array, return lower-triangular entries as a 1D array. This does not create a special Op but instead reduces to multiplication.

Parameters:: L (ArrayLike | InfixRV) – 2D square array with shape (N,N)
Returns:: 1D array of length N*(N+1)/2

Examples

>>> L = constant([[1.1,0],[2.2,3.3]])
>>> l = extract_tril(L)
>>> print_upstream(l)
shape  | statement
------ | ---------
(2, 2) | a = [[1.1 0. ] [2.2 3.3]]
(3,)   | b = [0 1 1]
(3,)   | c = [0 0 1]
(3,)   | d = vmap(index, [None, 0, 0], 3)(a,b,c)

pangolin.interface.print_upstream(*vars, **named_vars)[source]#

Prints all upstream variables in a friendly readable format.

Parameters:

vars (PyTree[RV]) – any number of pytrees containing RV
named_vars (RV) – single RV as keyword arguments, will be printed with those names

Examples

>>> r = RV(Constant(0.5))
>>> s = RV(Bernoulli(), r)
>>> t = RV(Constant(2))
>>> u = RV(Normal(), s, t)
>>> v = RV(Constant([75, 50, 99]))
>>> print_upstream([u, v]) # use autogenerated names
shape | statement
----- | ---------
()    | a = 0.5
()    | b ~ bernoulli(a)
()    | c = 2
()    | d ~ normal(b,c)
(3,)  | e = [75 50 99]
>>> print_upstream({'dog':[u], 'kitty':(v,)}) # any pytree is OK
shape | statement
----- | ---------
()    | a = 0.5
()    | b ~ bernoulli(a)
()    | c = 2
()    | d ~ normal(b,c)
(3,)  | e = [75 50 99]
>>> print_upstream(dog=u, kitty=v)
shape | statement
----- | ---------
()    | a = 0.5
()    | b ~ bernoulli(a)
()    | c = 2
()    | dog ~ normal(b,c)
(3,)  | kitty = [75 50 99]
>>> print_upstream(r=r, s=s, t=t, u=u, v=v) # control all names
shape | statement
----- | ---------
()    | r = 0.5
()    | s ~ bernoulli(r)
()    | t = 2
()    | u ~ normal(s,t)
(3,)  | v = [75 50 99]

pangolin.interface.RVLike#

A type class indicating either:

An RV
A NumPy array
Something that can be transformed into a NumPy array, such as a float or a list of lists of floats.

Many functions in this interface take RVLike arguments. If the argument to the function is not an RV, then an InfixRV with a Constant op is automatically created. So, for example, cos(2.5) is equivalent to

cos(InfixRV(pangolin.ir.Constant(2.5))).

Note that JAX arrays will typecheck as valid instances of RVLike (which is probably good) but so will strings (which is probably bad).

alias of ArrayLike | InfixRV

class pangolin.interface.Transform(forward, inverse, log_det_jac, n_biject_args=0)[source]#

Bases: object

Create a Transform object. A Transform can be used to wrap a function that produces a single RV with a random Op and turn it into a function the produces a single RV with type ir.Transformed.

Mathematically, the idea is that if P(X) is some density and Y=T(X) is a diffeomorphism, then P(Y=y) = P(X=T⁻¹(y)) × |det ∇T⁻¹(y)|. The goal here is that if you have some function that will produce the random variable X, this transform can be as a wrapper to create Y instead.

Three functions must be provided as inputs:

y = forward(x, b1, ..., bN) performs the forward transformation. This takes a single RVLike x along with some number of RVLike parameters b1, ..., bN and produces a single random InfixRV y as output. Any number of intermediate InfixRV can be created internally, but these must all be non-random. Must be a diffeomorphism with respect to x for any fixed values of b1, ..., bN.
x = inverse(y, b1, ..., bN) performs the inverse transformation, with similar properties as forward. Must be an inverse in the sense that inverse(forward(x,*args),*args)==x.
log_jac_det(x, y, b1, ..., bN) computes the log-determinant Jacobian log|det ∇forward(x,*args)| == -log|det ∇backwards(y,*args)|. Both x and y are provided since it may be more convenient to use one rather than the other.

When a transform is called, by default it returns a wrapped function that produces only a transformed random variable (y). However, if the keyword-only argument orig is True, then the wrapped function also returns a deterministic version of the original variable. That is, it returns x that is a deterministic function of y.

Parameters:

forward (Callable[[...], InfixRV]) – Function to compute forward transformation.
inverse (Callable[[...], InfixRV]) – Function to compute inverse transformation.
log_det_jac (Callable[[...], InfixRV]) – Function to compute log-determinant-Jacobian.
n_biject_args (int)

Returns:

Wrapper function that takes a function fun(*parents) -> InfixRV[Op] and returns a function wrapped(*args,*parents) -> InfixRV[Wrapped[Op,Bijector]]. (The bijector is constructed automatically.) The final Op (Wrapped[Op,Bijector]) is always random.

Examples

Create an inverse-exponential distributed random variable:

>>> reciprocal = Transform(lambda x: 1/x, lambda y: 1/y, lambda x, y: -2*log(abs(x)))
>>> print(exponential(1.5))
exponential(1.5)
>>> print(reciprocal(exponential)(1.5))
transformed(exponential,
            bijector(composite(1, [1, div], [[], [1, 0]]),
                     composite(1, [1, div], [[], [1, 0]]),
                     composite(2, [abs, log, -2, mul], [[0], [2], [], [4, 3]])))(1.5)

Create an inverse-exponential-distributed random variable. But also return a deterministic random variable where the inverse is un-done.

>>> reciprocal = Transform(lambda x: 1/x, lambda y: 1/y, lambda x, y: -2*log(abs(x)))
>>> y = reciprocal(exponential)(1.5)
>>> x = reciprocal.inverse(y)
>>> x.op
Div()
>>> y in x.parents
True

Do the same as the above in a single call.

>>> reciprocal = Transform(lambda x: 1/x, lambda y: 1/y, lambda x, y: -2*log(abs(x)))
>>> y, x = reciprocal.apply_and_invert(exponential)(1.5)
>>> x.op
Div()
>>> y in x.parents
True

Create a random variable distributed like x*4.4 for x ~ normal(3.3, 3.3)

>>> myfun = lambda a: normal(a,a)
>>> scale = Transform(lambda x, b: x*b, lambda y, b: y/b, lambda x, y, b: log(b), 1)
>>> x = myfun(3.3)
>>> print(x)
normal(3.3, 3.3)
>>> y = scale(myfun, 4.4)(3.3)
>>> print(y)
transformed(normal,
            bijector(mul,
                    div,
                    composite(3, [log], [[2]])), 1)(4.4, 3.3, 3.3)

Create a lognormal distribution parameterized in terms of the precision of the original normal.

>>> exp_tform = Transform(exp, log, lambda x,y: y)
>>> def normal_precision(mean, precision):
...     return normal(mean, 1/precision ** 0.5)
>>> x = exp_tform(normal_precision)(1.1, 5.5)
>>> print(x)
transformed(normal,
            bijector(exp,
                    log,
                    composite(2, [identity], [[1]])))(1.1, div(1, pow(5.5, 0.5)))

See also

pangolin.ir.RV

__getitem__(idx)[source]#

You can index an RV with the [] operator, e.g. as A[B,C].

Note that indexing with this interface is different (and simpler) than NumPy or JAX:

First, indexing is by default fully-orthogonal. This is done to avoid the utter insanity that is NumPy indexing with broadcasting, basic indexing, advanced indexing, and combinations of basic and advanced indexing. In this interface, if A, B, and C are RVs, then A[B,C].shape == A.shape + B.shape, and similarly if B or C are int / list of (list of) int / numpy array / slice.

Second, all axes must be indexed. For example, if A is a RV with 3 axes, then A[2] will trigger an exception. The idea of this is to make code more legible and self-enforcing. Instead you must write A[2, :, :] or A[2, ...].

Examples

>>> # Basic indexing
>>> A = constant([9,8,7,6,5,4])
>>> B = A[2]
>>> B.op
Index()
>>> B.parents[0] == A
True
>>> B.parents[1]
InfixRV(Constant(2))

>>> # indexing with a slice
>>> B = A[1::2]
>>> B.op
Index()
>>> B.parents[0] == A
True
>>> B.parents[1]
InfixRV(Constant([1,3,5]))

>>> # indexing with a combination of constants and slices
>>> A = constant([[3,4,5],[6,7,8]])
>>> B = A[[1,0],::2]
>>> B.op
Index()
>>> B.parents[0] == A
True
>>> B.parents[1]
InfixRV(Constant([1,0]))
>>> B.parents[2]
InfixRV(Constant([0,2]))

Parameters:: idx (_IdxType | tuple[_IdxType, ...])

property parent_ops#: Just a shortcut for tuple(p.op for p in self.parents). Intended mostly for testing.

property vindex#

Do vectorized indexing. That is, treat the RV as a scalar function from (scalar) indices to (scalar) output. Then, broadcast the indices according to standard broadcasting rules.

Technically rv.vindex returns a proxy object that can be indexed, so that rv.vindex[a,b,c] is equivalent to vindex(rv, a, b, c). Slices are converted to corresponding 1D arrays. Unless an ellipsis is given, the number of arguments must exactly match the number of dimensions of the array.

Examples

>>> A = constant([[0,1,2],[3,4,5]])
>>> B = A.vindex[[2,2,2,0,0,0],[0,1,2,0,1,2]]
>>> print_upstream(B)
shape  | statement
------ | ---------
(2, 3) | a = [[0 1 2] [3 4 5]]
(6,)   | b = [2 2 ... 0 0]
(6,)   | c = [0 1 ... 1 2]
(6,)   | d = vmap(index, [None, 0, 0], 6)(a,b,c)

class pangolin.interface.Broadcasting(*values)[source]#

Broadcasting behavior for scalar functions.

Used to change broadcasting behavior via config.

>>> config.broadcasting = Broadcasting.OFF    # no broadcasting
>>> config.broadcasting = Broadcasting.SIMPLE # simple broadcasting (default)
>>> config.broadcasting = Broadcasting.NUMPY  # numpy-style broadcasting

Broadcasting applies to “all-scalar” functions where all inputs and outputs are scalar. If broadcasting is set to OFF, all inputs must actually be scalar, and the output is scalar. If broadcasting is set to SIMPLE, then inputs can be non-scalar, but all non-scalar inputs must have exactly the same shape, which is the shape of the output If broadcasting is set to NUMPY, then inputs are broadcast against each other NumPy-style and the resulting shape is the shape of the output.

Suppose f(a,b,c) is a scalar function. Then here is how broadcasting behaves:

			Broadcasting mode
			OFF	SIMPLE	NUMPY
`a.shape`	`b.shape`	`c.shape`	`f(a,b,c).shape`
`()`	`()`	`()`	`()`	`()`	`()`
`(5,)`	`()`	`()`	n/a	`(5,)`	`(5,)`
`(5,)`	`(5,)`	`()`	n/a	`(5,)`	`(5,)`
`(5,)`	`(5,)`	`(5,)`	n/a	`(5,)`	`(5,)`
`(5,9)`	`()`	`()`	n/a	`(5,9)`	`(5,9)`
`(5,9)`	`(5,9)`	`()`	n/a	`(5,9)`	`(5,9)`
`(5,9)`	`(5,9)`	`(5,9)`	n/a	`(5,9)`	`(5,9)`
`(9,)`	`(5,9)`	`()`	n/a	n/a	`(5,9)`
`(5,)`	`(6,)`	`()`	n/a	n/a	n/a

This interface does not (currently) support automatic broadcasting for non-scalar functions.

OFF#: No broadcasting at all. Scalar functions only accept scalar arguments.

SIMPLE#: Simple broadcasting only. Arguments can be any shape, but non-scalar must have exactly the same shape.

NUMPY#: NumPy-style broadcasting. The only limitation is that broadcasting of singleton dimensions against non-singleton dimensions is not currently supported.

Examples

No broadcasting

>>> from pangolin import interface as pi
>>> pi.config.broadcasting = pi.Broadcasting.OFF
>>> pi.student_t(0,1,1).shape
()
>>> pi.student_t(0,np.ones(5),1).shape
Traceback (most recent call last):
    ...
ValueError: Non-scalar argument: [1. 1. 1. 1. 1.]

Simple broadcasting

>>> pi.config.broadcasting = pi.Broadcasting.SIMPLE
>>> pi.student_t(0,1,1).shape
()
>>> pi.student_t(0,np.ones(5),1).shape
(5,)
>>> pi.student_t(0,np.ones((5,3)),1).shape
(5, 3)
>>> pi.student_t(0,np.ones((5,3)),np.ones((5,3))).shape
(5, 3)
>>> pi.student_t(0,np.ones(3),np.ones((5,3))).shape
Traceback (most recent call last):
    ...
ValueError: Can't broadcast non-matching shapes (5, 3) and (3,)

NumPy broadcasting

>>> pi.config.broadcasting = pi.Broadcasting.NUMPY
>>> pi.student_t(0,1,1).shape
()
>>> pi.student_t(0,np.ones(5),1).shape
(5,)
>>> pi.student_t(0,np.ones((5,3)),1).shape
(5, 3)
>>> pi.student_t(0,np.ones((5,3)),np.ones((5,3))).shape
(5, 3)
>>> pi.student_t(0,np.ones(3),np.ones((5,3))).shape
(5, 3)

Restore to default

>>> pi.config.broadcasting = pi.Broadcasting.SIMPLE

You can also change this behavior temporarily in a code block by using the override context manager.

>>> from pangolin import interface as pi
>>> with pi.override(broadcasting="off"):
...     x = pi.normal(np.zeros(3), np.ones((5,3)))
Traceback (most recent call last):
...
ValueError: Non-scalar argument: [0. 0. 0.]
>>> with pi.override(broadcasting="numpy"):
...     x = pi.normal(np.zeros(3), np.ones((5,3)))
...     x.shape
(5, 3)

class pangolin.interface.Config(broadcasting=Broadcasting.SIMPLE)[source]#

Global configuration for Pangolin interface.

See also

override: Context manager for temporary config changes.
config: instance to use to actually make changes

Parameters:: broadcasting (Broadcasting)

pangolin.interface.config = Config(broadcasting=<Broadcasting.SIMPLE: 'simple'>)#: Singleton instance of Config class

pangolin.interface.override(**kwargs)[source]#

Temporarily override config values.

A context manager that sets config values for the duration of the block, then restores the original values on exit.

Parameters:

kwargs – Attribute names and their temporary values. See config for available options.

Raises:

AttributeError – If a key doesn’t match a config attribute.
ValueError – If a value is invalid for that config option.

Example

>>> from pangolin.interface import override, Broadcasting
>>> config.broadcasting
<Broadcasting.SIMPLE: 'simple'>
>>> with override(broadcasting="off"):
...     config.broadcasting
<Broadcasting.OFF: 'off'>
>>> config.broadcasting
<Broadcasting.SIMPLE: 'simple'>

pangolin.interface.Autoregressable = Autoregressable#

Type alias.

Type aliases are created through the type statement:

type Alias = int

In this example, Alias and int will be treated equivalently by static type checkers.

At runtime, Alias is an instance of TypeAliasType. The __name__ attribute holds the name of the type alias. The value of the type alias is stored in the __value__ attribute. It is evaluated lazily, so the value is computed only if the attribute is accessed.

Type aliases can also be generic:

type ListOrSet[T] = list[T] | set[T]

In this case, the type parameters of the alias are stored in the __type_params__ attribute.

See PEP 695 for more information.

pangolin.interface.Autoregressed = Autoregressed#

Type alias.

Type aliases are created through the type statement:

type Alias = int

In this example, Alias and int will be treated equivalently by static type checkers.

Type aliases can also be generic:

type ListOrSet[T] = list[T] | set[T]

In this case, the type parameters of the alias are stored in the __type_params__ attribute.

See PEP 695 for more information.

interface Module

Contents

interface Module#

Reference card#

Auto-casting#

Broadcasting#

API docs#

Sub-Packages#