liesel-devs / liesel

A probabilistic programming framework
https://liesel-project.org
MIT License
38 stars 2 forks source link

Make transform a method #174

Open jobrachem opened 7 months ago

jobrachem commented 7 months ago

This PR introduces a Var.transform method. Some notes:

Replacing GraphBuilder.transform

We currently have GraphBuilder.transform. The new method is in fact intended as a replacement for GraphBuilder.transform. Having Var.transform as a method on Var has the advantage that it is easier to find for users. Since the transformation is "doing something to a Var" and it does not actually require any functionality within the GraphBuilder, living as a method on a Var is a natural development for the transform method.

Behavior change

The method behaves similar to GraphBuilder.transform, with a few notable differences:

  1. It will not use the default event space bijector from tensorflow by default. This change is made to encourage users to either select their desired bijector manually, which is often sensible, or to request an automatic bijector manually. In both cases, users are more aware of what they are doing.
  2. The new method accepts bijector instances in addition to bijector classes. In fact, passing an instance is the preferred way of passing a bijector. Passing a bijector class is only supported if you actually defined *bijector_args or *bijector_kwargs to be passed to the bijector. This simplifies the code for the default case. More importantly, this fixes the graph representation after transformation, see below.

Using a bijector instance:

image

Compare to using the default event space bijector. Note that in this case, the same bijector is being used, but there are spurious edges from the nodes "v0" and "v1", the prior parameters, to the original variable "tau".

image

Deprecation and updated documentation

  1. I marked GraphBuilder.transform as deprecated and included directions towards the new method in its documentation.
  2. I updated the usage of the transformation method in the tutorials 01a-transform.md and 04-mcycle.md. As it turned out, usage in these tutorials was outdated anyway.

Notebook for testing

You can play around with the method in this notebook:

050-transform.ipynb.zip

Related issues

jobrachem commented 6 months ago

A log message on the default bijector being used would be nice.

jobrachem commented 6 months ago

@jobrachem will transfer the core benefits of the Var.transform method to the GraphBuilder.transform: 1) accepting a bijector instance, 2) logging the default event space bijector

jobrachem commented 3 months ago

@jobrachem will continue as follows:

  1. Implement Var.transform
  2. Make API between Var.transform and GraphBuilder.transform consistent
  3. Add a note to the docs of GraphBuilder.transform, referring to Var.transform