Refactor the user-facing looping API

[rlouf]

Let's open a discussion on a new API for the looping primitives in aesara.

Looping in aesara revolves around one Op, the Scan op. One can create a Scan Op by calling the function scan :

results, updates = aesara.scan(fn, sequences, outputs_info, non_sequences, n_steps)

With the following arguments:

• fn is the function that is called at each iteration;
• sequences is the list of sequences over which to iterate;
• outputs_info are the variables to be carried around;
• non_sequences is the list of arguments that are passed to fn at each step;
• n_steps (required) is the number of steps to take.

It returns two values:

• results a Variable or a list of Variable
• updates with the update rules for all the shared variables used in scan

Issues with the user interface

The user interface makes very little sense when you do not know what it is actually doing internally. Nevertheless users should be able to use the main features of aesara without having to know what happens in the background, the same way one does not have to know what tracing and XLA are all about when using JAX.

I have also identified the following difficulties with the user interface, most stem from the fact that the current implementation tries to do too many things with a single operator.

• Too much is infered from the type of output_info. It can be a list of Variable or adictionary that specifies the initial value and taps;
• One needs to refer to outputs_info to be able to make the distrinction between which returned values are carried over and which are just collected.
• The kind of control flow primitive Scan represents is determined by which arguments are provided;
• To make Scan behave as a while loop the body function needs to return an aesara.scan.utils.until object;
• When returning an aesara.utils.until instance to make Scan behave as a while loop the body function is executed at least once;
• The updates returned value is (unnecessarily?) cumbersome;
• We very often have to specify the arguments to scan (such as outputs_info, which is used very often) as a kwarg.
• scan only accepts TensorVariables as arguments. While the design choice simplifies many things, it pushes complexity out to downstream applications. This can be a real issue, see aehmc’s codebase.

I think it is possible to adress all these shortcomings with minimal change to the existing Scan operator. I suggest to provide user-facing APIs for condition- and collection-controlled loops.

Motus operandi

The idea is to deprecate the current functions and Ops progressively:

1. Add new user interface, points to existing Scan;
2. If the abstractions make sens, refactor Scan operator. Rename;
3. Make new user interface point to new Op;
4. Eventually deprecate both the old Op and user interface.

Proposal

Python has two different kinds of loop:

• Condition-controlled loops while
• Collection-controlled loops for x in iterable

and to reduce cognitive load we should provide aesara primitives that map to these constructions. I also suggest we try to simplify the interface as much as we can.

Automatic updates if possible

We should aim at completely removing the manual update logic.

Deprecate taps

Theano was used for deep learning, and implementing taps probably made sense if the group worked with RNNs a lot. I suggest we remove this feature: it is a nice to have, but a refactor should focus on what is essential. It should be easy for libraries that make heavy use of taps to implement a wrapper around the new scan primitive.

While loop

The first step is to implement an Op that maps to the standard while loop. We take inspiration from JAX and Tensorflow to propose the following:

final_value = aesara.while_loop(cond_fn, body_fn, init_value, n_steps)

With the following rough type signatures:

def body_fn(val: A) -> A
def cond_fn(val: A) -> bool
def while_loop(conf_fn, body_fn, init_value: A, n_steps: int) -> A

We could then replace the following example (taken from the documentation)

import aesara
import aesara.tensor as at


def power_of_2(previous_power, max_value):
    return previous_power*2, aesara.scan.utils.until(previous_power*2 > max_value)

max_value = at.scalar()
values, _ = aesara.scan(
    power_of_2,
    outputs_info = aet.constant(1.),
    non_sequences = max_value,
    n_steps = 1024
)

f = aesara.function([max_value], values)

print(f(45))

with the more explicit:

import aesara

max_value = at.scalar()

def cond_fn(previous_power):
    return previous_power * 2 > max_value

def body_fn(previous_power) :
    return previous_power ** 2

last_value = aesara.while_loop(cond_fn, body_fn, at.constant(1.), n_steps=1024)

f = aesara.function([max_value], last_value)

print(f(45))

With the following arguments to while_loop:

• cond_fn function such that we keep iterating while cond_fn(variables)=True;
• body_fn function called at each step body_fn(variables);
• init_vals the initial value(s) of the variable(s) either as a variable, a tuple or list (*args) or a dictionary (**kwargs);
• n_steps maximum number of steps (@brandonwillard can this be made optional, or does it have implications for e.g gradient computation?)

Comments:

• We do not collect the values by default, but users can implement collection within this design should they need to.
• cond_fn is checked before calling body_fn the first time;
• Updates should be performed automatically and the function does not return updates (@brandonwillard is that possible?)
• Output type and shape are the same at init_value’s

Scan

We take inspiration from JAX and Tensorflow implementation:

values = aesara.scan(body_fn, init, sequences=None, n_steps=None)

with the following rough type signature:

def body_fn(s: S, val: A) -> (A, B)
def scan(body_fn, init: A, sequences: Optional[List[S]], n_steps: Optional[int]) -> (A, List[B])

Here I don’t have enough experience with aesara to know if we need to provide the type of B via the user interface of if it can be infered. In the following I assume I can infer it.

Simple loop

We can implement a simple loop setting sequences to None and giving a value to n_step. So that:

import aesara
import aesara.tensor as at

k = at.iscalar("k")
A = at.vector("A")

values, updates = aesara.scan(fn=lambda prior_result, A: prior_result * A,
                              outputs_info=aet.ones_like(A),
                              non_sequences=A,
                              n_steps=k)
final_result = values[-1]

power = aesara.function(inputs=[A,k], outputs=final_result, updates=updates)

becomes:

import aesara
import aesara.tensor as at

k = at.iscalar("k")
A = at.vector("A")

def multiply(value):
    V, prior_result = value
    result = V * prior_result
    return (V, result), result

final_result, values = aesara.scan(multiply, (A, A), n_steps=k)

power = aesara.function(inputs=[A,k], outputs=final_result)

Iterate over the first dimension of a TensorVariable

import aesara
import aesara.tensor as at
import numpy

coefficients = at.vector("coefficients")
x = at.scalar("x")

max_coefficients_supported = 10000

def take_power(coefficient, power, free_variable):
    return coefficient * (free_variable ** power)

# Generate the components of the polynomial
components, updates = aesara.scan(fn=take_power,
                                  outputs_info=None,
                                  sequences=[coefficients, at.arange(max_coefficients_supported)],
                                  non_sequences=x)
# Sum them up
polynomial = components.sum()

# Compile a function
calculate_polynomial = aesara.function(inputs=[coefficients, x], outputs=polynomial)

Becomes:

import aesara
import aesara.tensor as at
import numpy

coefficients = at.vector("coefficients")
x = at.scalar("x")

max_coefficients_supported = 10000

def take_power(free_variable, s):
    coefficient, power = s
    return free_variable, coefficient * (free_variable ** power)

# Generate the components of the polynomial
_, values = aesara.scan(take_power, free_variable, (coefficients, at.arange(max_coefficients_supported)])
polynomial = values.sum()

# Compile a function
calculate_polynomial = aesara.function(inputs=[coefficients, x], outputs=polynomial)

[ricardoV94]

CC @junpenglao

[brandonwillard]

• n_steps maximum number of steps (@brandonwillard can this be made optional, or does it have implications for e.g gradient computation?)

From what I remember, n_steps is simply assumed to be available by the Scan implementation code (i.e. Scan.perform and the Cython implementation). My guess is that the code was initially designed for fixed-iteration-length, so that assumption was natural, then the while-looping feature was tacked on later without revising the previous assumptions (probably because refactoring that code is extremely cumbersome).

Aside from a guess at the "historical" reasons why n_step is required for a while-like Scan graph, I can't think of any reason why such a term would be necessary. It's not necessary at the symbolic graph representation level, nor is it at the computational level.

• Updates should be performed automatically and the function does not return updates (@brandonwillard is that possible?)

Using the auto-updates mechanism as an example, I don't see why it wouldn't be possible. We don't want the Variable objects themselves to hold this information, like they currently do with auto-updates, because our basic graph objects should soon become immutable; however, that's just an implementation detail.

We really need to do a full assessment of the updates feature, but, as I recall, all the updates information is essentially available/contained within the graphs output by Scan. I looked into somewhat recently this while working on AePPL, because I needed to generate the updates without using the scan interface function (see here).

The point is that we shouldn't need the caller to manage the updates, because that information is available to anyone who has the graph. If that's not actually the case, we need to make it so.

Aside from whether or not we can technically remove updates at the user-level, we need to consider the options/flexibility that's offered by the current approach and address that. For instance, if a user omits the updates generated by Scan when calling aesara.function, the shared variables involved will never change during each iteration. If this is something useful, or what a user wants, then we would need to add options for it in our fully-automated approach.

I don't think there needs to be much of a challenge involved in addressing that, though. For example, we could provide a keyword that specifies the shared variables to be updated, and our new interface would automatically handle the rest; otherwise, by default it would perform all updates.

After considering all this, I'm starting to think that removing updates from the Scan interface might not be all that difficult.

[rlouf]

After considering all this, I'm starting to think that removing updates from the Scan interface might not be all that difficult.

That's great news!

Aside from whether or not we can technically remove updates at the user-level, we need to consider the options/flexibility that's offered by the current approach and address that. For instance, if a user omits the updates generated by Scan when calling aesara.function, the shared variables involved will never change during each iteration. If this is something useful, or what a user wants, then we would need to add options for it in our fully-automated approach.

Regarding random variables, I think it would be less confusing, albeit slightly more verbose in many applications, to make users update the seed manually (with something similar to JAX's splitting mechanism). It is more explicit, and I find it easier to reason about the seed than about updates of shared variables.

[brandonwillard]

You're comparing two very distinct mechanisms: "splitting" and shared variables.

This "splitting" is just the idea of generating a new RNG state from an existing one (in this context, at least), and it's already deeply integrated into the RandomVariable interface, since the outputs of a RandomVariable node are another RandomState (i.e. the "split" state) and the sampled output.

Here's an example that avoids shared variables altogether:

import aesara
import aesara.tensor as at
import numpy as np
from aesara.tensor.random.type import random_generator_type

# A non-shared RNG state.  It will be an external input to the graph.
rng_at = random_generator_type()

X1_rv = at.random.normal(rng=rng_at)

# Get the "split" RNG object
new_rng = X1_rv.owner.outputs[0]

X2_rv = at.random.normal(rng=new_rng)

fn = aesara.function([rng_at], [X1_rv, X2_rv])


rng_val = np.random.default_rng(123)
fn(rng_val)
# [array(-0.98912135), array(-0.36778665)]
fn(rng_val)
# [array(-0.98912135), array(-0.36778665)]

# Change the RNG state
_ = rng_val.normal()

fn(rng_val)
# [array(-0.36778665), array(1.28792526)]
fn(rng_val)
# [array(-0.36778665), array(1.28792526)]

Shared variables are part of a higher level mechanism that introduces "stateful" terms to a graph. This statefulness works really well for RNG objects, because they're generally designed to be stateful and mutated in-place. More importantly, they can be used to produce much simpler graphs, because their use can remove the strong graph dependencies introduced by the approach above.

Notice how X2_rv is directly dependent on X1_rv:

aesara.dprint(fn)
# normal_rv{0, (0, 0), floatX, False}.1 [id A] ''   0
#  |<RandomGeneratorType> [id B]
#  |TensorConstant{[]} [id C]
#  |TensorConstant{11} [id D]
#  |TensorConstant{0.0} [id E]
#  |TensorConstant{1.0} [id F]
# normal_rv{0, (0, 0), floatX, True}.1 [id G] ''   1
#  |normal_rv{0, (0, 0), floatX, False}.0 [id A] ''   0
#  |TensorConstant{[]} [id C]
#  |TensorConstant{11} [id D]
#  |TensorConstant{0.0} [id E]
#  |TensorConstant{1.0} [id F]

That dependency is not exclusive to the RNG-object produced by X1_rv, but to the entirety of the X1_rv graph (i.e. its inputs and so forth). When using shared RNG variables, the two can be completely independent.

Remember, the way that this independence is accomplished is by creating independent shared RNG variables with different seeds. I believe we could do the same thing with non-shared RNG variables, we would just need to implement an RNG constructor Op that takes a seed argument. Lack of such an Op is probably one of the (historical) reasons for the reliance on shared variables, although another may be the challenges induced by Scan in this case.

If we started using non-shared RNG objects in the same way we're currently using the shared ones (i.e. constructing new ones with different seeds instead of chaining RNG state outputs as they're used), Scan still wouldn't work as expected unless we pass the RNG states around as "TAP" -1 values that are used by the next iteration. As I understand, the shared variable updates provide the per-iteration updates, but through a "statefulness" mechanism that's independent of Scan and RandomVariable.

Remember, Theano/Aesara graphs were designed to not carry "state", except through the shared variables/updates mechanisms. In part, this "statelessness" is only a conceptual ideal, and it can be broken by allowing in-place updates. There's an Op-level option for that and a rewrite that will enable it under certain circumstances (i.e. random_make_inplace). This inplace option was set to True for a while in Theano-PyMC, which is why my old articles seemed to operate so differently—and no longer work with the current codebase.

In other words, with an RNG state constructor Op and in-place updates set on RandomVariable Ops, this whole shared variable thing can be avoided completely—without unnecessarily complicating graphs with chained RNG-based dependencies; however, I'm not sure that we can simply enable in-place updates universally—especially not for shared variables.

Another option is to go with the chained situation initially and then strategically replace the unnecessary dependencies with in-place updates. This is the best sounding approach, but it should require some updates to random_make_inplace and possibly Scan. Plus, it would still require a "state" of some kind, like RandomStream. As a matter of fact, the whole approach could be immediately kick-started by changing RandomStream so that it no longer uses shared variables and it uses the RNG state outputs of its previous samples (e.g. the "split" RNG state produced by srng.normal() will be used as the rng value to the next srng.normal() call).

First and foremost, we need to preserve reproducibility of our graphs, and we can't mutate external inputs, so, for example, a NumPy RNG state object can't change after calling an Aesara compiled function that uses it as an argument. This shouldn't be a problem, though, because the arguments are generally copied by default, and use of an RNG constructor Op would avoid that situation entirely. Same with reproducibility: the graphs can't change in non-deterministic ways, so the same RNG state inputs should necessarily produce the same results for the same graphs.

Anyway, I'm pretty sure that these are the kinds of things we need to be considering in order to improve this situation.

[brandonwillard]

Here I don’t have enough experience with aesara to know if we need to provide the type of B via the user interface of if it can be infered. In the following I assume I can infer it.

The types of the outputs for the inner/loop function do not need to be specified separately. All one needs to do is construct dummy variables for the appropriate "inner-inputs" (i.e. the inputs to the inner function) and evaluate the inner function with those variables to get an "inner-graph" that represents the loop body. The output Variables of that inner-graph are the "inner-outputs". That's all that's needed for type inference in this case.

This is what the scan helper function does. It takes the "outer-input" sequences (i.e. the arrays over which we're iterating) and determines the Types of their iteration elements (i.e. the inner-inputs) so that it can construct the inner-graph as above and use its inner-outputs to determine the Types of the "outer-outputs" (i.e. the results produced by performing the iterations with the given loop body).

That's really all there is to Scan at the graph level. Unfortunately, the old/current implementation doesn't make that apparent; if anything, it hides that simplicity.

[brandonwillard]

We take inspiration from JAX and Tensorflow implementation:

values = aesara.scan(body_fn, init, sequences=None, n_steps=None)

with the following rough type signature:

def body_fn(s: S, val: A) -> (A, B)
def scan(body_fn, init: A, sequences: Optional[List[S]], n_steps: Optional[int]) -> (A, List[B

To clarify, the only substantial difference (in terms of the broad design) between the general "carry-tuple" approach used by JAX and the current Scan is that Scan uses a stricter, flattened version of the inner-function signature.

That was likely done to "simplify" the implementation and interpretation of Scan's internal elements, because—for example—lagged terms and their orders are explicitly given by the user. Also, there apparently wasn't established support for a tuple-like structure in Theano at the time.

Regardless, to determine that something is a "TAP" term according to Scan's nomenclature, simply observe the inner-function's input and output/carry connection pattern, which is pretty easy to do when you have a symbolic graph. This simple observation could've been used to remove/avoid the overly strict order-based constraints from the scan interface.

[rlouf]

Here is another possibility that fully embraces that fact that looping primitives are Ops. I kind of like it:

loop = aesara.while_loop(cond_fn, body_fn)
last_value = loop(init_value)

And

scan = aesara.scan(body_fn)
last_value, acc = scan(init_value, sequences)

[rlouf]

Again I need to check, but I think JAX doesn't accumulate because of constraints on the XLA side. If we don't have the same constraints then I don't see why we wouldn't accumulate. More use cases, API closer to scan's.

[rlouf]

Could accumulate be a boolean kwarg?

scan = aesara.scan(body_fn)
acc = scan(init_value, sequences, accumulate=True)

It could but I'd have to try it on a couple of mock examples. Definitely a property of the Op though.

[rlouf]

The easiest being I think to systematically carry over the first return value, accumulate the second. Easy to remember.

The more I see this API the more I like it.

[junpenglao]

acc = scan(init_value, sequences, accumulate: callable)

where accumulate = lambda i, cum_previous_x, current_x: ...

[junpenglao]

For none accumulate the operation is something like at.append or accumulate = lambda i, cum_previous_x, current_x: cum_previous_x.write(i, current_x) (which IIUC is the TF implementation. In TF scan is implemented using tf.while_loop)

[ricardoV94]

while_loop -> while

[rlouf]

Edit: while is a reserved word in python. This is probably JAX's while loop is also called while_loop.

[junpenglao]

yeah otherwise difficult to from xx import while

[brandonwillard]

I've just added two issues that, when resolved, would almost entirely remove the need to use shared variables when combining RandomVariable and Scan: #739 and #738.

With those in place (well, just #739, really), RandomStream could be made to work as it currently does without producing shared variables, and that would be all that's needed to make shared variables and updates entirely optional.