# Blockwise distribution of MultivariateNormal and TransformedDistribution

**URL:** <https://discuss.ai.google.dev/t/blockwise-distribution-of-multivariatenormal-and-transformeddistribution/31848>\
**Category:** General Discussion\
**Tags:** tf-probability, help\_request\
**Created:** [June 26, 2024, 10:09am UTC](https://discuss.ai.google.dev/t/blockwise-distribution-of-multivariatenormal-and-transformeddistribution/31848 "2024-06-26T10:09:29Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![doc\_56](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/doc_56/32/7725_2.png) [@doc\_56](https://discuss.ai.google.dev/u/doc_56)\
**Post date:** [June 26, 2024, 10:09am UTC](https://discuss.ai.google.dev/t/blockwise-distribution-of-multivariatenormal-and-transformeddistribution/31848/1 "2024-06-26T10:09:29Z")

</div>

Hi everyone,

I’m trying to use the` Blockwise` distribution in order to be able to sample from a `TransformedDistribution` and a `MultivariateNormal` and have a single vector as output. The code is the following:

```python
import tensorflow as tf
from tensorflow_probability import bijectors as tfb
from tensorflow_probability import distributions as tfd

mean = tf.random.normal(shape=(32, 256))
sigma = tf.random.normal(shape=(32, 256))
transformed_distribution = tfd.TransformedDistribution(
    distribution=tfd.Normal(loc=mean[:, 0], scale=sigma[:, 0]),
    bijector=tfb.BatchNormalization()
)
multivariate_normal = tfd.MultivariateNormalDiag(
    loc=mean[:, 1:], 
    scale_diag=sigma[:, 1:]
)
blockwise = tfd.Blockwise([transformed_distribution, multivariate_normal])
samples = blockwise.sample() # Here I expect samples.shape = (32, 256)

```

But It does not execute with success. Here’s the traceback:

> Traceback (most recent call last):  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/distributions/distribution.py”, line 1205, in sample  
> return self.\_call\_sample\_n(sample\_shape, seed, \*\*kwargs)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/distributions/distribution.py”, line 1182, in \_call\_sample\_n  
> samples = self.\_sample\_n(  
> ^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/distributions/blockwise.py”, line 350, in \_sample\_n  
> self.\_distribution.sample(n, seed=seed))  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/distributions/distribution.py”, line 1205, in sample  
> return self.\_call\_sample\_n(sample\_shape, seed, \*\*kwargs)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/distributions/joint\_distribution.py”, line 956, in \_call\_sample\_n  
> return self.\_sample\_n(  
> ^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/internal/distribution\_util.py”, line 1350, in \_fn  
> return fn(\*args, \*\*kwargs)  
> ^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/distributions/joint\_distribution.py”, line 699, in \_sample\_n  
> xs = self.\_call\_execute\_model(  
> ^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/distributions/joint\_distribution.py”, line 850, in \_call\_execute\_model  
> return self.\_execute\_model(  
> ^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/distributions/joint\_distribution.py”, line 1005, in \_execute\_model  
> next\_value, traced\_values = sample\_and\_trace\_fn(  
> ^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/distributions/joint\_distribution.py”, line 132, in trace\_values\_only  
> ret = trace\_distributions\_and\_values(dist, sample\_shape, seed, value)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/distributions/joint\_distribution.py”, line 115, in trace\_distributions\_and\_values  
> value = dist.sample(sample\_shape, seed=seed)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/distributions/distribution.py”, line 1205, in sample  
> return self.\_call\_sample\_n(sample\_shape, seed, \*\*kwargs)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/distributions/transformed\_distribution.py”, line 338, in \_call\_sample\_n  
> return self.bijector.forward(x, \*\*bijector\_kwargs)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/bijectors/bijector.py”, line 1326, in forward  
> return self.\_call\_forward(x, name, \*\*kwargs)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/bijectors/bijector.py”, line 1308, in \_call\_forward  
> return self.\_cache.forward(x, \*\*kwargs)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/internal/cache\_util.py”, line 334, in forward  
> return self.\_lookup(x, self.\_forward\_name, self.\_inverse\_name, \*\*kwargs)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/internal/cache\_util.py”, line 493, in \_lookup  
> self.\_invoke(input, forward\_name, kwargs, attrs))  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/internal/cache\_util.py”, line 532, in \_invoke  
> return getattr(self.bijector, fn\_name)(input, \*\*kwargs)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/bijectors/batch\_normalization.py”, line 224, in \_forward  
> return self.\_de\_normalize(x)  
> ^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/bijectors/batch\_normalization.py”, line 216, in \_de\_normalize  
> mean = broadcast\_fn(self.batchnorm.moving\_mean)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow\_probability/python/bijectors/batch\_normalization.py”, line 204, in \_broadcast  
> return tf.reshape(v, broadcast\_shape)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow/python/ops/weak\_tensor\_ops.py”, line 88, in wrapper  
> return op(\*args, \*\*kwargs)  
> ^^^^^^^^^^^^^^^^^^^  
> File “/.venv/lib/python3.11/site-packages/tensorflow/python/util/traceback\_utils.py”, line 153, in error\_handler  
> raise e.with\_traceback(filtered\_tb) from None  
> File “/.venv/lib/python3.11/site-packages/tensorflow/python/eager/execute.py”, line 53, in quick\_execute  
> tensors = pywrap\_tfe.TFE\_Py\_Execute(ctx.\_handle, device\_name, op\_name,  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> tensorflow.python.framework.errors\_impl.InvalidArgumentError:  
> {{function\_node \__wrapped\_\_Reshape\_device_/job:localhost/replica:0/task:0/device:GPU:0}}  
> Input to reshape is a tensor with 32 values, but the requested shape has 1 [Op:Reshape]  
> python-BaseException

I assume it has to do with the way I’m building the `Blockwise` distribution and the batch size and event size of its components but I’m not able to fix it.

Thanks for your help.