Fix MultiVariateNormalDiag and similar distributions leaking tracers.

[rainx0r]

Currently, MultiVariateNormalDiag and other such distributions that depend on the DiagLinear and DiagPlusLowRankLinear distributions are not jittable (see #230) because they leak tracers.

This means that they cannot be returned from jitted functions. This is a bug because these distributions and bijectors all inherit from Jittable and yet are not able to be used as inputs or outputs of such functions as intended.

The issue is that these bijectors attempt to "inherit" from their internal _bijector objects by doing the following in their __init__ method:

self.forward = self._bijector.forward

However, I found that this actually creates a closure over self._bijector at trace time which in turn leaks the tracer held within self._bijector. Simply implementing the corresponding functions that at call-time access self._bijector fixes the issue.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[llandsmeer]

Hi, I can confirm the PR solved the issue I was running into (leaked tracers)