Replies: 4 comments 12 replies
-
|
I think the multiple-output/default output issue should be treated distinctly from the seeding issue. It doesn't make sense to refactor the way seeding works just because we are hiding one of the outputs by default and that causes problems (we could always decide to not hide). |
Beta Was this translation helpful? Give feedback.
-
The multiple-output issue is completely unrelated and discussed in #1036, we just need to sketch a way to solve it at the representation level. (Well that's not completely true because it creeps in with the multiple output of Of course we can return all the outputs, but that doesn't solve much. The default output issue should be understood as an example of how convenient design choices that were made at the RNG level are bothering us in seemingly unrelated places. The point of this is not so much about solving issues as it is about focusing on the representation itself, hoping that a "correct" representation would not cause this kind of issue in the first place. |
Beta Was this translation helpful? Give feedback.
-
I believe this was mentioned before, but the basic interface and functionality of Unless that's not completely true and I've missed something important, this proposal is almost entirely about adopting JAX's high-level interface and refactoring Speaking of copying the RNG state: one of the reasons |
Beta Was this translation helpful? Give feedback.
-
My main concern about import aesara.tensor as at
srng = at.random.RandomStream(0)
X_at = first_fun(srng)
Y_at = second_fun(srng)Also, unless I've missed something, some work needs to be done to expose this "advancing" mechanism in the IR so we can move forward with aesara-devs/aemcmc#66.
It is about managing the PRNG state more explicitly at the graph level, without threading it through otherwise independent portions of the graph.
But assuming the rng = np.random.Generator(np.random.Philox(key=k, counter=c))
rng.normal()Although it is not shown here, most of the work is done in |
Beta Was this translation helpful? Give feedback.
-
|
It's also important to note that none of this is either new or JAX-specific. The literature on this design long predates JAX. |
Beta Was this translation helpful? Give feedback.
-
Unless
The same could be done just as easily with only Currently, RNG state advancement is represented implicitly in IR by In other words, the
As I mentioned above, it looks like |
Beta Was this translation helpful? Give feedback.
-
Yes; I was only referring to the explicit changes implied. Regardless, these are my first impressions, and I still need to wrap my head around all this. |
Beta Was this translation helpful? Give feedback.
-
Then how can we use the same
We need to distinguish between the representation of the RNG state in Aesara's iR and any particular implementation, and we can think about those independently as long as there's a simple mapping from the former to the latter. That being said, it looks like def split(seed_seq):
return seed_seq.spawn(2)Indeed, if
It is indeed the underlying question.
Which leads to threading of otherwise independent portions of the code, which makes graph introspection and manipulation difficult.
I am not sure it removes the need for IR-level state advancements. For users it probably does in most cases (although not all of them, cf above), but representing this explicitly in the IR would greatly help with some rewrites down the line.
In the proposed design nodes are not directly inter-dependent; there is a clean separation between these nodes' subgraph and the RNG states' subgraph.
I think the proposed
It looks like it, and I am happy to write less code and just use |
Beta Was this translation helpful? Give feedback.
-
|
@rlouf How would seeding update happen across function calls (or within Scan iterations) with your approach? Edit: Also thanks for the detailed write up! |
Beta Was this translation helpful? Give feedback.
-
|
I don't know if it's feasible to update a seedsequence / spawn multiple generators from it every iteration. Could be a costly process, but maybe it's not? |
Beta Was this translation helpful? Give feedback.
-
Yes, it most likely is. I've entertained the idea of seeding every draw with something like To provide some context, here's an example of the underlying objects involved: import numpy as np
rng = np.random.default_rng(293)
rng._bit_generator.state
# {'bit_generator': 'PCG64',
# 'state': {'state': 88109460180158103940050617480421485605,
# 'inc': 200169897845647114283224307892191878717},
# 'has_uint32': 0,
# 'uinteger': 0}
rng._bit_generator._seed_seq.state
# {'entropy': 293, 'spawn_key': (), 'pool_size': 4, 'n_children_spawned': 0}
rng.normal()
# -0.581704568682811
rng._bit_generator.state
# {'bit_generator': 'PCG64',
# 'state': {'state': 281649526231773309250369509310953017910,
# 'inc': 200169897845647114283224307892191878717},
# 'has_uint32': 0,
# 'uinteger': 0}The following is how the rng._bit_generator._seed_seq.generate_state(4)
# array([3622451066, 357958987, 3911150219, 525595238], dtype=uint32)Those values are then used to set the We need to compare the seeding/ Using The only way forward I can imagine that uses a "seeding every draw" approach would almost certainly involve some lowish-level work that—for example—circumvents the standard interfaces, or uses a very specific PRNG. In all other cases, and as we currently do, we need to (implicitly) rely on In summary:
I'm really hoping that I'm wrong about the inefficiency of seeding every draw, though. |
Beta Was this translation helpful? Give feedback.
-
|
Yeah, and even if we stuck to Philox, are we even sure that gen = np.random.Generator(np.random.Philox(counter=rng.counter, key=rng.key))
gen.normal(0, 1)
new_rng, _ = at.random.split(rng)
gen = np.random.Generator(np.random.Philox(counter=rng.counter, key=rng.key))
gen.normal(0, 1)would be as fast as rng = np.random.default_rng()
rng.normal(0, 1)
rng.normal(0, 1)? We would need to run some benchmarks with this, and But what if there were a way to trade the seeding cost at compile time for a rewrite? We could for instance easily rewrite the following graph: rng = at.random.RandState(0)
rng_x, rng_y, rng_z = at.random.splitn(rng, 3)
x_rv = at.random.normal(rng_x, 0, 1)
y_rv = at.random.normal(rng_y, 0, 1)
z_rv = at.random.normal(rng_z, x_rv, y_rv)Which can be visualised as
To this one in NumPy: rng = np.random.default_rng(0)
x_rv = rng.normal(0, 1)
y_rv = rng.normal(0, 1)
z_rv = rng.normal(x_rv, y_rv)
Which would be as fast as it gets at runtime. I am not sure what the tradeoffs of this solution would be besides implementing the rewrites, but, conceptually, it certainly moves towards Python/NumPy being a backend just like Numba or JAX. |
Beta Was this translation helpful? Give feedback.
-
That should be what the |
Beta Was this translation helpful? Give feedback.
-
|
Indeed, this example was maybe too simple. What I sketched could only happen if we rewrote the graph with the new So, assuming seeding at every draw is very expensive for PRNGs that are not counter-based (but we should benchmark this), we currently have 3 options:
Warning: the following is highly speculative, might be completely stupid (I'll write some code and report back) I am wondering if there might be a 4th solution, using the srng = at.random.RandomStream(0)
srng_1, srng_2 = at.random.split(srng)
x_rv = call(srng_1)
w_rv = call(srng_1) # uses the same sequence
y_rv = call(srng_2) # different sequence
z_rv = srng.normal(0,1) # advances the original streamI have no idea if that's possible, but I find it's a much better interface than the constant splitting happening in JAX, while giving the user more control over the rng state. The drawback is that we would be threading draws that are sequential, even though RVs have no data dependence. |
Beta Was this translation helpful? Give feedback.
-
|
Following up after a while... This redesign is definitely worthwhile as a means of decoupling RNG state evolution from the process of drawing samples from RNG states, which is valuable at an interface/representation level, so I'm willing to go forward with it for that reason alone. |
Beta Was this translation helpful? Give feedback.
-
Summary
I propose a new approach to represent
RandomVariableand PRNG states in Aesara’s IR, based on the design of splittable PRNGs. The representation introduces minimal change to the existingRandomVariableinterface while being more expressive. It should be easy to transpile to Aesara’s current compilation target, and is compatible with the higher-levelRandomStreaminterface.References
On counter-based PRNGs:
On splittable PRNG design for functional programs:
On the NumPy side:
Desiderata
A good PRNG design satisfies the following conditions:
Motivation behind this proposal
This proposal is motivated by two issues that illustrate the shortcomings of the current representation of
RandomVariable\s and PRNG states in Aesara:In #1036 the use of
default_outputto hide the PRNG state from the user is causing multiple headaches in the etuplization ofRandomVariable\s and unification/reification of expressions withRandomVariable\s. This is the onlyOpin Aesara that makes use of this property, and to “special-casing” the etuplization logic forRandomVariable\s often appeared as the easiest solution.In #66 in AeMCMC, expressing expansions like the following convolution of two normal variables is overly complex:
It is indeed not clear what the values of
rng_xandrng_yshould be given the value ofrng. A few other application-related shortcomings of the current representation will be given below.Proposal
In the following we focus on the symbolic representation of random variables and PRNG states in Aesara’s IR. We leave discussions about compilation targets and solution to the previous issues for the end.
If we represent the internal state of the PRNG by the type
RandState(short forRandomStateType), the current design ofRandomVariable\s can be summarized by the following simplified signature:In other words,
RandomVariable\s are responsible for both advancing the state of the PRNG, and producing a random value. This double responsibility is what creates graph dependencies between nodes that have otherwise no data dependency. The following snippet illustrates this:As we can see in the graph representation,
rng_x(id C) is being used as an input toyandrng_y(id I) is being used as an input toz. There is however no data dependency betweenx,yorz. The intuition that they should not be linked is probably what led to “hiding” these PRNG state outputs so they are not re-used, and theRandomStreaminterface.Creating spurious sequential dependencies by threading PRNG states is indeed unsatisfactory from a representation perspective, and unnecessarily complicates the rewrites. It is also problematic for two other reasons:
A natural idea is to simplify the design of
RandomVariable\s so that it is only responsible for one thing: generating a random value from a PRNG state. TheOpthus creates anApplynode that takes aRandState(using the above notation) as input and outputs a (random)Variable:Providing a
RandStateto aRandomVariableneeds to intentional, and this must be reflected in the user interface. We thus makerngan explicit input of theRandomVariable’s__call__method. This way a user can write:Or, if they want the PRNG state to be shared (silly example, but a legitimate need):
This interface presupposes the existence of two operators. First, to build reproducible programs, we need an operator that creates a
RandStatefrom a seed, which can be the constructor ofRandStateitself:And then, we need another operator that creates an updated
RandomStatefrom aRandomState, so thatRandomVariable\s created with these two different states would output different numbers. Let’s call itnext:We can thus fill in the blanks in the previous code examples:
The code has been specifically formatted to illustrate what we gain from this approach.
x_rv,y_rvandz_rvhave lost their direct dependency; we could easily execute these three statements in parallel. What we have done implicitly is to create two graphs: the graph between random variables which reflects the dependencies (or lack thereof) on each other’s values, and the graph of the updates of the PRNG states. These graphs almost evolve in parallel.This is similat to what I understand the
RandomStreaminterface does: moving the updates of the PRNG states to theupdategraphs generated by Aesara’s shared variables.The
nextoperator is however not completely satisfactory. Let us consider a more complex situation, wherecallis a function that requires aRandomState:We can easily find an implementation of
callthat makes the previous code generate a random state collision:To avoid this kind of issues, we must thus require user-defined functions to return the last PRNG state along the result:
Threading PRNG state is still necessary to guarantee correctness and the two
callfunctions cannot be called in parallel. The issue arises because, even though we have separated PRNG state update and random value generation, our symbolic structure is still sequential: eachRandStatehas one and only one ancestor. We can of course circumvent this issue knowing how many timesnextis called within the function, by “jumping” the same number of times to obtainrng_y, but this can quickly become complex (what ifcallis imported from somewhere else?).It would make things easier if a
RandState\s could have several children, and if each of these child led to separate streams of random number. Let us define the followingsplitoperator:We require that we can never get the same
RandStateby callingsplitany number of times on either the left or right returned state. In other words,splitshould implicitly defines a binary tree in which all the nodes are unique. This can be easily represented by lettingRandStateholding a number in binary format. The leftmost child state is obtained by appending0to the parent’s state and the rightmost child state by appending1:If the generator called by
RandomVariablecan be made a deterministic function of this binary value, the computations are fully reproducible. We added akeyattribute that can be specified by the user at initialization to seed the PRNG state. The tree structure is of course explicit in our graph representation, sincelandrdepend onrngvia thesplitoperator. Nevertheless, we can increment this internal state when building the graph in a way that allows us to compile without traversing the graph.The
nextoperator we previously defined becomes redundant within this representation. Since its interaction with thesplitoperator would require careful thought we leave it aside in the following. Using the new operator our toy example becomes:Note that the “main” sub-graph that contains random variables, and the PRNG sub-graph are still minimally connected.
Finally, it is also natural to implement the
splitnoperator represented by:So we can write the following code:
Implementation
When it comes to practical implementations, this representation is only convenient for counter-based PRNGs like
Philoximplemented in NumPy: we generate a pair of(key, counter)from ourRandState\s and pass these as an input to the generator.RandStateandsplitimplementationThe mock implementation of$\mathcal{S}$ of real PRNGs does not usually extend indefinitely. In practice we will need to compress the state using a hashing function that also increments the
RandStateandsplitabove is naive in the sense that the counter spacekey. To be immediately compatible with NumPy in theperformfunction we can use Philox’s hash function to update the state as we build the graph. Since the hash is deterministic we can still walk theRandStatetree in our representation and cheaply recompute the states should we need to.Op and Variable implementations to come.
RandomVariableThe modifications to
RandomVariableOps are minimal:__call__now takes aRandStateas a positional argument;make_nodeonly returnsout_var. Thedefault_outputattribute is not needed anymore.RandomStreamWe can keep the
RandomStreamAPI, use a shared variable to hold theRandStateand handle the splitting internally. The RNG sub-graphs are now found in the updates’ graph.In a second time we may consider instantiating
RandStateas shared variables by default to decouple both the random variable and the PRNG state graphs. I am not sure of the tradeoffs here, but it may alleviate concerns related to graph rewrites.Compilation
It is essential that our representation of PRNG states and
RandomVariable\s in the graph can be easily transpiled to the existing targets (C, Numba, JAX) and future targets. In the following I outline the transpilation process for the current targets.Numba
After #1245 Aesara will support NumPy’s Generator API. Furthermore NumPy has support for Philox as a BitGenerator, a counter-based PRNG which can easily accomodate splittable PRNG representations. Assuming we can map each path in the PRNG graph to a
(key, counter)tuple, the transpilation ofRandomStream\s using the PhiloxBitGeneratorshould be straighforward. For the explicit splitting interface, we can directly translate theRandomVariable\s to NumPyGenerator\s and seed these generators at compile time. So that:Becomes:
JAX
Transpilation to JAX would be straightforward, as JAX uses a splittable PRNG representation. We will simply need to perform the following substitutions:
Back to the motivating issues
The problems linked to the existence of the
default_outputattribute disappear sinceRandomVariable\s do not return PRNG states anymore. The one-to-many difficulty we are facing with the relations between etuplized graphs also disappears with asplitoperator. Using the example from the beginning we can for instance write:Which is guaranteed to be collision-free by construction of the
splitoperator, as long as the PRNG state used by the original normal distribution isn’t passed to asplitoperator somewhere else in the original graph (todo: specify API requirements to guarantee uniqueness of the random numbers).Related
Related, not necessarily a justification for this approach, discussions across the Ae* projects:
OpFromGraphand updatesBeta Was this translation helpful? Give feedback.
All reactions