ir Module

Contents

ir Module#

The pangolin IR for defining joint distributions over collections of random variables. Note that while it is certainly possible to define probabilistic model directly in terms of this IR, it is quite a “low-level” notation, and end-users of Pangolin are not typically expected to manipulate this IR directly. Instead, users would typically use an pangolin.interface that provides a friendlier way of creating these models.

This IR only represents groups of dependent random variables. All actual functionality is left to inference engines.

The two core abstractions in the IR are the Op and the RV. An Op represents a simple function or conditional distribution, while an RV consists of a single Op and a (possibly empty) tuple of parent Op. Thus, the IR is essentially just a directed graph where each node is an RV with a single Op.

That’s the whole IR. Unlike in Pyro / NumPyro / PyMC / Tensorflow Probability, Random variables do not have names, and there is no notion of a “model” class. In Pangolin, you just work with RVs. You can assign those RVs to named variables or put them into tuples or lists or dicts if you want. But that’s up to you.

Other notes:

  • All RV and Op are immutable. Nothing is allowed to change after initialization.

  • All RV have static shapes. The shape of an RV is recursively determined by its Op and the shapes of its parent RV.

Types of Op:

Type

Ops

Constants

Constant

Arithmetic

Add Sub Mul Div

Trigonometry

Arccos Arccosh Arcsin Arcsinh Arctan Arctanh Cos Cosh Sin Sinh Tan Tanh

Other scalar functions

Pow Abs Exp InvLogit Log Loggamma Logit Step

Linear algebra

Matmul Transpose Inv Cholesky

Other multivariate functions

Sum Softmax Diag DiagMatrix

Scalar distributions

Normal NormalPrec Lognormal Cauchy Bernoulli BernoulliLogit Beta Binomial Categorical Uniform BetaBinomial Exponential Gamma Poisson StudentT

Multivariate distributions

MultiNormal Multinomial Dirichlet Wishart

Control flow

VMap Composite Autoregressive

Indexing

Index

In addition, this module provides print_upstream, which provides a nice human-readable description of an entire graph.

pangolin.ir.Shape#

A Shape is just a tuple of ints

alias of tuple[int, …]

class pangolin.ir.Op[source]#

Bases: ABC

Abstract base class for operators. An Op represents a deterministic function or distribution. All functionality for sampling or evaluating densities is left to inference backends. An Op is frozen after initialization.

Simple Op typically have no parameters. For example, Add() just represents scalar addition. The Op itself does not indicate what variables are being added together. Similarly, Normal() just represents a normal distribution, but it does not itself indicate what the mean and scale of that distribution are. More complex classes have constructors that take any parameters that must be fixed in order to maintain fixed shapes. For example Sum takes the axis to sum over.

A concrete Op class must provide a _random class attribute, indicating if an op is a conditional distribution (True) or a deterministic function (False). For classes where this is fixed, this can just be a bool. For classes where this depends on the particular instance of the class (e.g. VMap) this can be a function taking the op instance and returning a bool. This property is accessed by random. In addition, the docstring for random in subclasses is automatically overridden based on _random. (If _random is bool, the fixed value is given as a string. If _random is a function, the docstring for _random is copied to random in the subclass.)

A concrete Op class must also provide a function _get_shape as a class attribute, which will be called by get_shape. This is a function, not a constant, because some Op (e.g. MultiNormal or Matmul) may have differnet shapes depending on the shapes of the inputs. This _get_shape function should take the shapes of all parents and compute the shape of the output. It is also expected that the _get_shape method will do error checking—e.g. verify that the correct number of parents are provided and the shapes of the parents are coherent with each other. This method is called by RV at construction to ensure that the shapes of all random variables are coherent.

An Op must provide an __eq__ method that indicates mathematical equality. Thus, if op1 = Normal() and op2 = Normal() then, op1 == op2, even though op1 and op2 are distinct objects. However, if op3 = Sum(0) and op4 = Sum(1) then op3 != op4. This base class provides a simple implementation that checks if the two arguments have the same type. This must be overridden for classes like Sum that have constructors that take parameters. (In those cases, __hash__ must also be overridden.)

property random: bool#

Is this class a distribution (True) or a deterministic function (False)?

get_shape(*parents_shapes)[source]#

Given the shapes of parents, what is the shape of the output of this Op?

Parameters:

parents_shapes (Shape) – Shape of each parent

Returns:

shape of output of this op

Raises:

  • TypeError – If an incorrect number of parents are provided.

  • ValueError – If the shapes of the parents are incoherent.

Return type:

Shape

__eq__(other)[source]#

Are self and other mathematically equal?

Parameters:

other – op to compare to.

Return type:

bool

property name: str#

Returns the name of the op class as a string

__str__()[source]#

Provides a more compact representation, e.g. normal instead of Normal()

Return type:

str

class pangolin.ir.Identity[source]#

Bases: Op

Represents an “identity” Op. Has one parent.

property random: bool#

False

class pangolin.ir.Constant(value)[source]#

Bases: Op

Represents a “constant” Op. Has no parents. Data is always stored as a numpy array. If you want to live dangerously, you may be able to switch to jax’s version of numpy by setting ir.np = jax.numpy.

Parameters:

value (ArrayLike) – Some constant value that is either a numpy array or something that can be cast to a numpy array.

value#

The actual stored data, stored as an immutable numpy array

__eq__(other)[source]#

Returns True if other is of type Constant and other.value is exactly the same as self.value.

Parameters:

other (Op)

Return type:

bool

get_shape()#

If len(parents_shapes)>0, raises ValueError. Otherwise, returns the shape of value.

Return type:

Shape

property random: bool#

False

class pangolin.ir.ScalarOp[source]#

Bases: Op, ABC

Abstract class to conveniently create “all-scalar” ops. These are simple ops that:

  1. Have class-fixed randomness.

  2. Have a class-fixed number of parents.

  3. Expect all parents to be scalar.

  4. Output a scalar.

To create a concrete instance of this class, all that’s necessary is to specify two things:

  1. _random a bool indicating if the Op is a conditional distribution or a deterministic function. Unlike in a general Op this cannot be a function.

  2. _expected_parents which can either be the number of parents (an int) or a dictionary mapping parent names to descriptions. If a dictionary is given, those parent names will be included in the documentation, and may be used when constructing interface functions.

You can also optionally provide _wikipedia string. This is only used to create a link in the documentation. If not provided, there is a heuristic to try to guess a link for distributions.

class pangolin.ir.Add[source]#

Bases: ScalarOp

Represents a Add Op. Takes no parameters. When used in an RV, expects 2 scalar parent(s):

  • a: parent 0

  • b: parent 1

Examples

>>> op = Add()
>>> op
Add()
>>> print(op)
add
>>> op.random
False
>>> op.get_shape((), ())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Sub[source]#

Bases: ScalarOp

Represents a Sub Op. Takes no parameters. When used in an RV, expects 2 scalar parent(s):

  • a: parent 0

  • b: parent 1

Examples

>>> op = Sub()
>>> op
Sub()
>>> print(op)
sub
>>> op.random
False
>>> op.get_shape((), ())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Mul[source]#

Bases: ScalarOp

Represents a Mul Op. Takes no parameters. When used in an RV, expects 2 scalar parent(s):

  • a: parent 0

  • b: parent 1

Examples

>>> op = Mul()
>>> op
Mul()
>>> print(op)
mul
>>> op.random
False
>>> op.get_shape((), ())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Div[source]#

Bases: ScalarOp

Represents a Div Op. Takes no parameters. When used in an RV, expects 2 scalar parent(s):

  • a: parent 0

  • b: parent 1

Examples

>>> op = Div()
>>> op
Div()
>>> print(op)
div
>>> op.random
False
>>> op.get_shape((), ())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Pow[source]#

Bases: ScalarOp

Represents a Pow Op. Takes no parameters. When used in an RV, expects 2 scalar parent(s):

  • a: parent 0

  • b: parent 1

Examples

>>> op = Pow()
>>> op
Pow()
>>> print(op)
pow
>>> op.random
False
>>> op.get_shape((), ())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Arccos[source]#

Bases: ScalarOp

Represents a Arccos Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Arccos()
>>> op
Arccos()
>>> print(op)
arccos
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Arccosh[source]#

Bases: ScalarOp

Represents a Arccosh Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Arccosh()
>>> op
Arccosh()
>>> print(op)
arccosh
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Arcsin[source]#

Bases: ScalarOp

Represents a Arcsin Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Arcsin()
>>> op
Arcsin()
>>> print(op)
arcsin
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Arcsinh[source]#

Bases: ScalarOp

Represents a Arcsinh Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Arcsinh()
>>> op
Arcsinh()
>>> print(op)
arcsinh
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Arctan[source]#

Bases: ScalarOp

Represents a Arctan Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Arctan()
>>> op
Arctan()
>>> print(op)
arctan
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Arctanh[source]#

Bases: ScalarOp

Represents a Arctanh Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Arctanh()
>>> op
Arctanh()
>>> print(op)
arctanh
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Cos[source]#

Bases: ScalarOp

Represents a Cos Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Cos()
>>> op
Cos()
>>> print(op)
cos
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Cosh[source]#

Bases: ScalarOp

Represents a Cosh Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Cosh()
>>> op
Cosh()
>>> print(op)
cosh
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Sin[source]#

Bases: ScalarOp

Represents a Sin Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Sin()
>>> op
Sin()
>>> print(op)
sin
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Sinh[source]#

Bases: ScalarOp

Represents a Sinh Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Sinh()
>>> op
Sinh()
>>> print(op)
sinh
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Tan[source]#

Bases: ScalarOp

Represents a Tan Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Tan()
>>> op
Tan()
>>> print(op)
tan
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Tanh[source]#

Bases: ScalarOp

Represents a Tanh Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Tanh()
>>> op
Tanh()
>>> print(op)
tanh
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Abs[source]#

Bases: ScalarOp

Represents a Abs Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Abs()
>>> op
Abs()
>>> print(op)
abs
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Exp[source]#

Bases: ScalarOp

Represents a Exp Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Exp()
>>> op
Exp()
>>> print(op)
exp
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.InvLogit[source]#

Bases: ScalarOp

Represents a InvLogit Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = InvLogit()
>>> op
InvLogit()
>>> print(op)
inv_logit
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Log[source]#

Bases: ScalarOp

Represents a Log Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Log()
>>> op
Log()
>>> print(op)
log
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Logit[source]#

Bases: ScalarOp

Represents a Logit Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Logit()
>>> op
Logit()
>>> print(op)
logit
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Step[source]#

Bases: ScalarOp

Represents a Step Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Step()
>>> op
Step()
>>> print(op)
step
>>> op.random
False
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Loggamma[source]#

Bases: ScalarOp

Represents a Loggamma Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • a: parent 0

Examples

>>> op = Loggamma()
>>> op
Loggamma()
>>> print(op)
loggamma
>>> op.random
False
>>> op.get_shape(())
()

Notes

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

get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

False

class pangolin.ir.Matmul[source]#

Bases: Op

A class that does matrix multiplication, following the rules of numpy.matmul. Currently only 1d and 2d arrays are supported.

Examples

>>> op = Matmul()
>>> op
Matmul()
>>> print(op)
matmul
>>> op.random
False
get_shape(a_shape, b_shape)#

Get the shape of applying a matmul to given shapes.

Parameters:

  • a_shape (Shape) – shape of first argument

  • b_shape (Shape) – shape of second argument

Return type:

Shape

Examples

>>> op = Matmul()
>>> op.get_shape((4,), (4,)) # inner product
()
>>> op.get_shape((5,4), (4,)) # matrix-vector
(5,)
>>> op.get_shape((5,), (5,4)) # vector-matrix
(4,)
>>> op.get_shape((5,4), (4,3)) # matrix-matrix
(5, 3)
>>> op.get_shape((5,5), (10, 10)) # incoherent shapes, error expected
Traceback (most recent call last):
...
ValueError: Matmul parent shapes do not match ((5, 5) vs. (10, 10))
property random: bool#

False

class pangolin.ir.Inv[source]#

Bases: Op

Take the inverse of a square matrix

get_shape(p_shape)#
Parameters:

p_shape (Shape) – A square 2D shape

Return type:

Shape

Returns

Same as p_shape

property random: bool#

False

class pangolin.ir.Cholesky[source]#

Bases: Op

Do a cholesky decomposition

get_shape(p_shape)#
Parameters:

p_shape (Shape) – A square 2D shape

Return type:

Shape

Returns

Same as p_shape

property random: bool#

False

class pangolin.ir.Transpose[source]#

Bases: Op

Take a transpose.

Unlike in many libraries this only works for 2D arrays, rather than reversing all dimensions. The logic is that 99% of the time, calling transpose on a 1D array and getting the original array (like in NumPy / JAX / Pytorch) is a bug. And calling transpose on a 4d array and reversing the dimensions is almost never useful. So, 2D transpose only!

get_shape(p_shape)#
Parameters:

p_shape (Shape) – A 2D shape

Return type:

Shape

Returns

p_shape, except reversed

property random: bool#

False

class pangolin.ir.Diag[source]#

Bases: Op

Extract the diagaonal of a square matrix. (Use DiagMatrix to construct.)

get_shape(p_shape)#
Parameters:

p_shape (Shape) – A 2D shape

Return type:

Shape

Returns

tuple with one element equal to p_shape[0] and p_shape[1] (or error if different)

property random: bool#

False

class pangolin.ir.DiagMatrix[source]#

Bases: Op

Extract the diagaonal of a square matrix. (Use DiagMatrix to construct.)

get_shape(p_shape)#
Parameters:

p_shape (Shape) – A 1D shape

Return type:

Shape

Returns

tuple with two elements both equal to p_shape[0]

property random: bool#

False

class pangolin.ir.Softmax[source]#

Bases: Op

property random: bool#

False

class pangolin.ir.Sum(axis)[source]#

Bases: Op

Create a Sum instance :param axis: What axis to sum over. :type axis: int

property random: bool#

False

class pangolin.ir.Normal[source]#

Bases: ScalarOp

Represents a Normal Op. Takes no parameters. When used in an RV, expects 2 scalar parent(s):

  • mu: location / mean

  • sigma: scale / standard deviation

Examples

>>> op = Normal()
>>> op
Normal()
>>> print(op)
normal
>>> op.random
True
>>> op.get_shape((), ())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

True

class pangolin.ir.NormalPrec[source]#

Bases: ScalarOp

Represents a NormalPrec Op. Takes no parameters. When used in an RV, expects 2 scalar parent(s):

  • mu: location / mean

  • tau: precision / inverse variance

Examples

>>> op = NormalPrec()
>>> op
NormalPrec()
>>> print(op)
normal_prec
>>> op.random
True
>>> op.get_shape((), ())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

True

class pangolin.ir.Lognormal[source]#

Bases: ScalarOp

Represents a Lognormal Op. Takes no parameters. When used in an RV, expects 2 scalar parent(s):

  • mu: logarithm of location

  • sigma: logarithm of scale (not sigma squared!)

Examples

>>> op = Lognormal()
>>> op
Lognormal()
>>> print(op)
lognormal
>>> op.random
True
>>> op.get_shape((), ())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

True

class pangolin.ir.Bernoulli[source]#

Bases: ScalarOp

Represents a Bernoulli Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • theta: probability (between 0 and 1)

Examples

>>> op = Bernoulli()
>>> op
Bernoulli()
>>> print(op)
bernoulli
>>> op.random
True
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

True

class pangolin.ir.BernoulliLogit[source]#

Bases: ScalarOp

Represents a BernoulliLogit Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • theta: logit of probability (unbounded)

Examples

>>> op = BernoulliLogit()
>>> op
BernoulliLogit()
>>> print(op)
bernoulli_logit
>>> op.random
True
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

True

class pangolin.ir.Binomial[source]#

Bases: ScalarOp

Represents a Binomial Op. Takes no parameters. When used in an RV, expects 2 scalar parent(s):

  • N: number of trials

  • theta: probability of success for each trial

Examples

>>> op = Binomial()
>>> op
Binomial()
>>> print(op)
binomial
>>> op.random
True
>>> op.get_shape((), ())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

True

class pangolin.ir.Cauchy[source]#

Bases: ScalarOp

Represents a Cauchy Op. Takes no parameters. When used in an RV, expects 2 scalar parent(s):

  • mu: location

  • sigma: scale

Examples

>>> op = Cauchy()
>>> op
Cauchy()
>>> print(op)
cauchy
>>> op.random
True
>>> op.get_shape((), ())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

True

class pangolin.ir.Uniform[source]#

Bases: ScalarOp

Represents a Uniform Op. Takes no parameters. When used in an RV, expects 2 scalar parent(s):

  • alpha: lower bound

  • beta: upper bound

Examples

>>> op = Uniform()
>>> op
Uniform()
>>> print(op)
uniform
>>> op.random
True
>>> op.get_shape((), ())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

True

class pangolin.ir.Beta[source]#

Bases: ScalarOp

Represents a Beta Op. Takes no parameters. When used in an RV, expects 2 scalar parent(s):

  • alpha: shape

  • beta: shape

Examples

>>> op = Beta()
>>> op
Beta()
>>> print(op)
beta
>>> op.random
True
>>> op.get_shape((), ())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

True

class pangolin.ir.Exponential[source]#

Bases: ScalarOp

Represents a Exponential Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • beta: rate / inverse scale

Examples

>>> op = Exponential()
>>> op
Exponential()
>>> print(op)
exponential
>>> op.random
True
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

True

class pangolin.ir.Gamma[source]#

Bases: ScalarOp

Represents a Gamma Op. Takes no parameters. When used in an RV, expects 2 scalar parent(s):

  • alpha: shape

  • beta: rate / inverse scale

Examples

>>> op = Gamma()
>>> op
Gamma()
>>> print(op)
gamma
>>> op.random
True
>>> op.get_shape((), ())
()

Notes

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

get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

True

class pangolin.ir.Poisson[source]#

Bases: ScalarOp

Represents a Poisson Op. Takes no parameters. When used in an RV, expects 1 scalar parent(s):

  • lambd: lambda

Examples

>>> op = Poisson()
>>> op
Poisson()
>>> print(op)
poisson
>>> op.random
True
>>> op.get_shape(())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

True

class pangolin.ir.BetaBinomial[source]#

Bases: ScalarOp

Represents a BetaBinomial Op. Takes no parameters. When used in an RV, expects 3 scalar parent(s):

  • N: as in binomial dist

  • alpha: as in beta dist

  • beta: as in beta dist

Examples

>>> op = BetaBinomial()
>>> op
BetaBinomial()
>>> print(op)
beta_binomial
>>> op.random
True
>>> op.get_shape((), (), ())
()

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.

get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

True

class pangolin.ir.StudentT[source]#

Bases: ScalarOp

Represents a StudentT Op. Takes no parameters. When used in an RV, expects 3 scalar parent(s):

  • nu: degress of freedom

  • mu: location (often 0)

  • sigma: scale (often 1)

Examples

>>> op = StudentT()
>>> op
StudentT()
>>> print(op)
student_t
>>> op.random
True
>>> op.get_shape((), (), ())
()
get_shape(*parents_shapes)#

Checks that correct number of parents are given and each has shape (). Always returns ().

Parameters:

parents_shapes (Shape) – Shape of each parent, all must be ().

Returns:

()

Raises:

  • TypeError – Incorrect number of parents

  • ValueError – Parent shapes not all ().

Return type:

Shape

property random: bool#

True

class pangolin.ir.MultiNormal[source]#

Bases: Op

Create a MultiNormal distribution. Takes no parameters.

When used in an RV, expects two parents: The mean and covariance.

get_shape(vec_shape, mat_shape)#

Checks that first argument is a 1D vector and second argument is a square 2D array with sizes the same as the first argument.

Parameters:
Return type:

Shape

property random: bool#

True

class pangolin.ir.Categorical[source]#

Bases: Op

Create a Categorical distribution. Takes no parameters.

When used in an RV, expects one parents: A 1-D vector of weights.

get_shape(weights_shape)#
Parameters:

weights_shape (Shape) – Should be 1D.

Returns:

()

Return type:

Shape

property random: bool#

True

class pangolin.ir.Multinomial[source]#

Bases: Op

Create a Multinomial Op. Takes no parameters.

When used in an RV, a multinomial op expects a first parent n (a scalar) indicating the number of observations and a second parent p indicating a vector of probabilities (a 1D array). Note that this is different from Stan (which doesn’t need n to be passed).

get_shape(n_shape, p_shape)#
Parameters:

  • n_shape (Shape) – Must be ()

  • p_shape (Shape) – Must be 1D

Returns:

Same as p_shape

Return type:

Shape

property random: bool#

True

class pangolin.ir.Dirichlet[source]#

Bases: Op

Create a Dirichlet Op. Takes no parameters.

When used in an RV, a Dirichlet op expects one parent, namely the concentration

get_shape(concentration_shape)#
Parameters:

concentration_shape (Shape) – 1D vector

Returns:

1D vector, same as concentration_shape.

Return type:

Shape

property random: bool#

True

class pangolin.ir.Wishart[source]#

Bases: Op

Create a Wishart op. Takes no parameters.

When used in an RV, expects two parameters: nu (degrees of freedom, scalar) and S (symmetric positive-definite scale matrix)

get_shape(nu_shape, S_shape)#
Parameters:

  • nu_shape (Shape) – must be ()

  • S_shape (Shape) – must be 2d square array shape

Returns:

Same as S_shape

Return type:

Shape

property random: bool#

True

class pangolin.ir.VMap(base_op, in_axes, axis_size=None)[source]#

Bases: Op, Generic

Create a VMap Op. That’s one specific op vectorized over some number of arguments.

All arguments here are heavily inspired by jax.lax.vmap although note that VMap only maps a single Op. (The vmap function in the interface elsewhere takes an arbitrary function and transforms it into a graph of RVs with VMap Op .)

Parameters:

  • base_op (BaseOp) – The Op to be mapped

  • in_axes (Sequence[int | None]) – What axis to map for each argument of the op (each can be a non-negative int or None)

  • axis_size (int | None) – The size of the mapped axis/axes. Optional unless all elements of in_axes are None.

Examples

>>> # diagonal normal
>>> op = VMap(Normal(), in_axes=[0, 0])
>>> op.get_shape((5,), (5,))
(5,)

>>> # diagonal normal with shared scale
>>> op = VMap(Normal(), in_axes=[0, None])
>>> op.get_shape((5,), ())
(5,)

>>> # 2D diagonal normal with shared scale
>>> op = VMap(VMap(Normal(), in_axes=[0, None]), in_axes=[0, None])
>>> op.get_shape((6,4), ())
(6, 4)

>>> # 2D diagonal normal with "outer product" of location and scale
>>> op = VMap(VMap(Normal(), in_axes=[None, 0]), in_axes=[0, None])
>>> op.get_shape((6,), (4,))
(6, 4)

>>> # 2D diagonal normal with "outer product" in the other order
>>> op = VMap(VMap(Normal(), in_axes=[0, None]), in_axes=[None, 0])
>>> op.get_shape((6,), (4,))
(4, 6)
__str__()[source]#

Return a string representation of the VMap op. Just like __repr__ except uses str for calling the recursive distribution.

get_shape(*parents_shapes)#

Expects shapes corresponding to the shapes expected by base_op but with extra axes as dictated by axis_size.

Parameters:

parents_shapes (Shape)

Return type:

Shape

property random: bool#

Equal to base_op.random

class pangolin.ir.Composite(num_inputs, ops, par_nums)[source]#

Bases: Op, Generic

A composite distribution can have any number of inputs. par_nums counts all the inputs first (in order) followed by the variables generated by the composite distribution itself.

Parameters:

  • num_inputs (int) – The number of input RV.

  • ops (tuple[*tuple[Op, ...], LastOp]) – tuple of Op to call (type system will insist this is a tuple, although using another sequence works).

  • par_nums (Sequence[Sequence[int]]) – Sequence of Sequence of ints indicating argument position for each Op.

property random: bool#

Equal to ops[-1].random

class pangolin.ir.Autoregressive(base_op, length, in_axes, where_self=0)[source]#

Bases: Op, Generic

Represents an autoregressive distribution

Parameters:

  • base_op (O)

  • length (int)

  • in_axes (Sequence[int | None])

  • where_self (int)

property random: bool#

Equal to self.base_op.random

class pangolin.ir.Index[source]#

Bases: Op

Create an Op to index into a RV in a simple orthogonal way. Does not deal with slices. Does not deal with broadcasting. The number of indices must always be the same as the number of dimensions of the RV. These things can be supported by the interface.

The shape of the output is always the concatenation of the shapes of the indices.

Examples

>>> op = Index()
>>> op.get_shape((5,5,5),(3,),(2,7,9),())
(3, 2, 7, 9)
property random: bool#

False

pangolin.ir.index_orthogonal(array, *index_arrays)[source]#

Create orthogonal index arrays for advanced indexing.

pangolin.ir.index_orthogonal_no_slices(A, *index_arrays)[source]#

Create orthogonal index arrays for advanced indexing.

class pangolin.ir.Bijector(forward, inverse, log_det_jac)[source]#

Bases: Op

The intended use for bijectors is for defining transformed distributions where we start with the density P(x) and we would like to quickly evaluate P(y) where y is some transformed version of x.

In general if Y=T(X) then

P_X(x) = P_Y(T(x)) × |det ∇T(x)|

and

P_Y(y) = P_X(T⁻¹(y)) × |det ∇T⁻¹(y)|.

That is, P_Y is the pushforward of P_X and P_X is the pullback of P_Y.

We assume that this distribution will generally be used as a “pullback”, meaning that an op is known for P(X) and one wishes to define an op for P(Y). In order to implement log_prob` for ``P(Y) it will be necessary to evaluate T⁻¹(y) and |det ∇T⁻¹(y)|. In order to sample P(Y) it will be necessary to evaluate T(x).

Parameters:

  • forward (Op) – Implements T(x)

  • inverse (Op) – Implements T⁻¹(y)

  • log_det_jac (Op) – Implements log(|det ∇T(x)|) == -log(|det ∇T⁻¹(y)|) (given both x and y).

Examples

An implementation of the exp transformation.

>>> forward = Exp()
>>> inverse = Log()
>>> log_jac_det = Composite(
...     2,
...     (Identity(),),
...     [[0]])
>>> bijector = Bijector(forward, inverse, log_jac_det)
>>> print(bijector)
bijector(exp, log, composite(2, [identity], [[0]]))
property random: bool#

False

class pangolin.ir.ExpBijector[source]#

Bases: NamedBijector

The Exp bijector. This is given as a subclass mostly to simplify printing.

property random: bool#

False

class pangolin.ir.Transformed(base_op, bijector, n_biject_args=0)[source]#

Bases: Op, Generic

Given some distribution P(X) and some bijection Y=T(X), create the distribution P(Y).

Transformed(base_op, transform_op, num_base_args, transformed_arg) Represents the result of drawing a sample from base_op and then putting that sample into a given position in transform_op.

If you do this:

>>> bijector = Bijector(forward, inverse, inverse_log_jac_det)
>>> op = Transformed(base_op, bijector, n_biject_args)
>>> y = RV(op, *p)

That produces a RV y that is equal in distribution to doing this:

>>> p_biject = p[:n_biject_args]
>>> p_base   = p[n_biject_args:]
>>> x = RV(base_op, *p_base)
>>> y = RV(forward, *p_biject)

However, in the former case y is considered random.

Examples

This is equivalent to a lognormal:

>>> op = Transformed(Normal(), ExpBijector())
>>> op
Transformed(Normal(), ExpBijector())
>>> print(op)
transformed(normal, exp_bijector)
>>> op.get_shape((),())
()
Parameters:

  • base_op (O) – The base op to transform (must be random)

  • bijection – The bijection to apply

  • n_biject_args (int) – Number of parameter parents for bijection (default 0)

  • bijector (B)

__str__()[source]#

Return a string representation of the VMap op. Just like __repr__ except uses str for calling the recursive distribution.

property random: bool#

True

class pangolin.ir.RV(op, *parents)[source]#

Bases: Node, Generic

A RV is essentially just an Op and a tuple of parent RV.

Parameters:

  • op (Op) – The Op defining the RV

  • *parents – The parents of the RV

Examples

>>> constant_op = Constant(3)
>>> x = RV(constant_op)
>>> x
RV(Constant(3))
>>> x.op
Constant(3)
>>> x.parents
()
>>> y = RV(constant_op)
>>> y
RV(Constant(3))
>>> normal_op = Normal()
>>> z = RV(normal_op, x, y)
>>> z
RV(Normal(), RV(Constant(3)), RV(Constant(3)))
>>> z.parents[0] == x
True
>>> z.parents[1] == y
True
property shape: Shape#

The shape of the RV. (A tuple of ints.)

property ndim: int#

The number of dimensions of the RV. Equal to the length of shape.

__setattr__(key, value)[source]#

Set attribute. Special case to freeze after init.

pangolin.ir.rv_equal(A, B)[source]#

Are A and B distributionally equal? That is, are they guaranteed to always have the same value? This is defined by the following rules:

  1. If A.op is random, then A is is equal to B if and only if they refer to the same object in memory.

  2. If A.op is non-random, then A is equal to B if and only if they have the same Op and their parents are equal (defined recursively using this function).

This function is implemented using caching. This doesn’t change the results since RV and Op s are immutable.

Parameters:

  • A (RV) – the first RV to be compared

  • B (RV) – the second RV to be compared

Return type:

bool

Examples

Constants behave as you’d expect.

>>> a = RV(Constant(0.5))
>>> b = RV(Constant(0.5))
>>> c = RV(Constant(0.7))
>>> # same object always equal
>>> rv_equal(a, a)
True
>>> # equivalent non-random ops, all (nonexistent) parents equal
>>> rv_equal(a, b)
True
>>> # non-equivalent ops always non-equal
>>> rv_equal(a, c)
False

Random Op are always non-equal unless they are same object.

>>> rv_equal(RV(Bernoulli(), a), RV(Bernoulli(), a))
False
>>> rv_equal(RV(Bernoulli(), a), RV(Bernoulli(), b))
False
>>> rv_equal(RV(Bernoulli(), a), RV(Bernoulli(), c))
False

Non-random op are if ops are equal and parents are recursively equal.

>>> rv_equal(RV(Exp(), a), RV(Exp(), a))
True
>>> rv_equal(RV(Exp(), a), RV(Exp(), b))
True
>>> rv_equal(RV(Exp(), a), RV(Exp(), c))
False
pangolin.ir.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]
Contents