View source on GitHub
|
Perform a scan with an associative binary operation, in parallel.
tfp.substrates.jax.math.scan_associative(
fn, elems, max_num_levels=48, axis=0, validate_args=False, name=None
)
The associative scan operation computes the cumulative sum, or
all-prefix sum, of a set of
elements under an associative binary operation [1]. For example, using the
ordinary addition operator fn = lambda a, b: a + b, this is equivalent to
the ordinary cumulative sum tf.math.cumsum along axis 0. This method
supports the general case of arbitrary associative binary operations operating
on Tensors or structures of Tensors:
scan_associative(fn, elems) = tf.stack([
elems[0],
fn(elems[0], elems[1]),
fn(elems[0], fn(elems[1], elems[2])),
...
fn(elems[0], fn(elems[1], fn(..., fn(elems[-2], elems[-1]))),
], axis=0)
The associative structure allows the computation to be decomposed
and executed by parallel reduction. Where a naive sequential
implementation would loop over all N elements, this method requires
only a logarithmic number (2 * ceil(log_2 N)) of sequential steps, and
can thus yield substantial performance speedups from hardware-accelerated
vectorization. The total number of invocations of the binary operation
(including those performed in parallel) is
2 * (N / 2 + N / 4 + ... + 1) = 2N - 2
--- i.e., approximately twice as many as a naive approach.
[1] Blelloch, Guy E. Prefix sums and their applications Technical Report CMU-CS-90-190, School of Computer Science, Carnegie Mellon University, 1990.
Args | |
|---|---|
fn
|
Python callable implementing an associative binary operation with
signature r = fn(a, b). This must satisfy associativity:
fn(a, fn(b, c)) == fn(fn(a, b), c). The inputs and result are
(possibly nested structures of) Tensor(s), matching elems. Each
Tensor has a batch dimension in place of elem_length; the fn
is expected to map over this dimension. The result r has the same shape
(and structure) as the two inputs a and b.
|
elems
|
A (possibly nested structure of) Tensor(s), each with dimension
elem_length along axis. Note that elem_length determines the number
of recursive steps required to perform the scan: if, in graph mode,
this is not statically available, then ops will be created to
handle any elem_length up to the maximum dimension of a Tensor.
|
max_num_levels
|
Python int. The axis of the tensors in elems must have
size less than 2**(max_num_levels + 1). The default value is
sufficiently large for most needs. Lowering this value can reduce
graph-building time when scan_associative is used with inputs of unknown
shape.
Default value: 48.
|
axis
|
Tensor int axis along which to perform the scan.
|
validate_args
|
Python bool. When True, runtime checks
for invalid inputs are performed. This may carry a performance cost.
Default value: False.
|
name
|
Python str name prefixed to ops created by this function.
|
Returns | |
|---|---|
result
|
A (possibly nested structure of) Tensor(s) of the same shape
and structure as elems, in which the kth element is the result of
recursively applying fn to combine the first k elements of
elems. For example, given elems = [a, b, c, ...], the result
would be [a, fn(a, b), fn(fn(a, b), c), ...].
|
Examples
from tensorflow_probability.python.internal.backend import jax as tf
import tensorflow_probability as tfp; tfp = tfp.substrates.jax
import operator
# Example 1: Partials sums of numbers.
tfp.math.scan_associative(operator.add, tf.range(0, 4))
# ==> [ 0, 1, 3, 6]
# Example 2: Partial products of random matrices.
dist = tfp.distributions.Normal(loc=0., scale=1.)
matrices = dist.sample(sample_shape=[100, 2, 2])
tfp.math.scan_associative(tf.matmul, matrices)