-
-
Notifications
You must be signed in to change notification settings - Fork 45
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[WIP] Composable transforms #947
Conversation
Codegen produces sensible outputs for single transforms. Still a good many TODOs/problems before chained transforms would compile
Can you show me an example of this? Or do you mean users literally call a function to check this? tmk in the compiler only transformed parameters do these checks where constraints don't actually apply but just checks. In that case wouldn't we just call multiple checks? |
Does data not get checked as input? This was a comment @bob-carpenter made to me when discussing it in front of a whiteboard, so he may be able to weigh in better. My understanding is that checking if something meets multiple constraint requirements would require doing the freeing transform in the middle. For example, a hypothetical value declared with constraints |
In the blocks The constraints on the The problems arise when applying constraints to already constrained values. So if I write
then what do I know about So if we want to match what happens with parameter transforms when checking bounds in the other blocks, we'd have to unwind the transforms step-by-step and check each step. That is, for
we want to validate as follows.
Alternatively, I think it'd make sense to only allow a single constraint on variables other than parameters. |
If we were only interested in allowing multiple transforms on parameters, I could probably tackle this with a slightly different strategy that might make some stuff simpler, but it would very tightly couple that idea (read: if someone later decided that was a poor choice and wanted MT other places, it would be as difficult as the first time around). The other option is maintain my current approach but just assert that the data block (..etc) cannot have anything other than a primitive (single) transform. That means the debug data generation and constraint check system etc could all remain mostly unchanged, and the only open issue to resolve here would be the dimension issues. Thoughts? I'm personally happy to put a nice error message inside debug data gen and constraint checking such that it gets punted to if/when someone decides later to implement that. Then when crafting the parser I can either disallow multiple constraints for non-parameters, or catch it during the semantic check. |
I'm trying to think of ways multiple transforms in a block other than |
May be incorrect
Still need to work on type param for read call
I'm good with
This starts delving into feature creep a bit, but I'd like to think about how user supplied constraints would interact here. I think what we are saying here is all well and good as long as in the future for constraining transformed parameters we do something where the main jist is to just expose the matrix[N, M] X = lb_constrain(X_other); We generate something like Eigen::Matrix<local_scalar_t__, -1, -1> X = lb_constrain<jacobian__>(X_other, lp); And then I think in |
That is how I would imagine it yes. Nothing I'm doing should make that any harder, and if anything might make it a little easier. The biggest issue from that standpoint I think is for transforms that change the dimensionality (like the matrix constraints we've discussed over email, or If you want to look at the current output I think I've nested everything as calls correctly, rather than repeated assignments. I still need to change the template parameters in the |
In general, a transform consists of several parts:
We use (1) for transforming inits and (2)+(3) for transforming parameters in the There's no reason in principle these can't be chained if the dimensions conform, but if you want to keep track of the ranges beyond the last transform, that's much trickier. The real question I have is whether we should try to encapsulate (1)-(4) into what TensorFlow Probability calls a "bijector" (that's not quite right as our unit vectors don't use a bijective transform). |
Currently getting CI failures about template substitution See Error:
|
Okay at least some of them are because I need to figure out how to pass the dimension args to the cholesky constraints. Not sure why |
Oh, looking those over we need the |
A little cleaner
I think I have a decent abstraction for FnReadParam for all of the transforms we currently implement, but I'm baking in some pretty strong assumptions by doing so that will have to change if we let users define arbitrary constraints (particularly if they change dimension) I'm starting to wonder if there is merit toward reintroducing an internal FnConstrain which handles that sort of stuff. I know #872 only just got rid of that. |
Okay, so I think there are 3 main things needed before the compiler work here is done.
Once these are done we can start worrying about how to actually instantiate this/the language design for it. I think I'll want at least one working version before this gets merged so we can test some things, even if it ends up in a separate PR. |
I just wanted to ping @rybern on this -- this is the WIP PR for composable transforms that touches a lot of that codegen. We should chat about how that could affect the tuples PR |
I think we've decided #971 supersedes this PR. |
Overview
This PR refactors our internal representation of transformations to include a concept of 'chaining' several in a row. The idea is roughly based on TFP's Chain bijector.
The code generation is based on turning something which currently would be codgen'd as
into
with the obvious extension that you can add additional reassignments to add additional constraints, e.g.
The same is done (in reverse order) for freeing.
There are several areas of the code which are not yet completed (hence the WIP). They are all marked with the string "TR TODO" in either a comment or a failwith statement.
Remaining work here
Transform_Mir
,Mir_Utils
, andAst_to_Mir
makes assumptions about sizes/dimensions based on the transform associated with the declaration, which still needs to be expanded or rethoughtFuture work
As a follow on PR, the parser will be updated to parse declarations like
As brought up in #659. Further extensions to allow more general compositions are left for others to consider.
Submission Checklist
Release notes
Internal refactor to allow composing multiple transforms of a single variable
Copyright and Licensing
By submitting this pull request, the copyright holder is agreeing to
license the submitted work under the BSD 3-clause license (https://opensource.org/licenses/BSD-3-Clause)