Generating ande valuating 200 Normals and reducing them

Question

I am trying to estimate a normal density using a quadratic approximation in tensorflow (code 4.14 from McElreath's Statistical Rethinking).

The code I have so far is:

import pandas as pd
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from  tensorflow_probability import distributions as tfd

_BASE_URL = "https://raw.githubusercontent.com/rmcelreath/rethinking/Experimental/data"

 HOWELL_DATASET_PATH = f"{_BASE_URL}/Howell1.csv"

df = pd.read_csv(HOWELL_DATASET_PATH, sep=';')
df = df[df['age'] >= 18]

mu = tf.linspace(start=140.0, stop=160.0, num=200)
sigma= tf.linspace(start=4.0, stop=9.0, num=200)

tf.reduce_sum(tfd.Normal(loc=mu, scale=sigma).log_prob(df.height))

This fails due to df having shape (352,) whilst I am creating (200,) points for my normal distribution to be evaluated on.

However

tf.reduce_sum(tfd.Normal(loc=mu, scale=sigma).log_prob(2))

and

tf.reduce_sum(tfd.Normal(loc=mu[0], scale=sigma[0]).log_prob(df.height))

both work.

I need to create a (200, 352) tensor - one Normal for each mu, sigma on my grid, and then evaluate it with my sample data - df. The question I have is: how do I do this?

Answer 1

So, I figured out that one way to do it would be to create a (200, 200, 352) grid and then reshape, and the rest of the calculations follow straightforwardly.

import pandas as pd
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from  tensorflow_probability import distributions as tfd

_BASE_URL = "https://raw.githubusercontent.com/rmcelreath/rethinking/Experimental/data"

 HOWELL_DATASET_PATH = f"{_BASE_URL}/Howell1.csv"

df = pd.read_csv(HOWELL_DATASET_PATH, sep=';')
df = df[df['age'] >= 18]


mu = tf.linspace(start=140.0, stop=160.0, num=200)
sigma = tf.linspace(start=7.0, stop=9.0, num=200)

means, variances, _  = tf.meshgrid(mu, sigma,  np.zeros((352,)).astype(np.float32))
means = tf.reshape(means, [40000, 352])
variances = tf.reshape(variances, [40000, 352])

normal = tfd.Normal(loc=means, scale=variances)

log_lik = tf.reduce_sum(normal.log_prob(df.height), axis=1)

logprob_mu = tfd.Normal(178.0, 20.0).log_prob(means)
logprob_sigma = tfd.Uniform(low=0.0, high=50.0).log_prob(variances)

log_joint_prod = log_lik + logprob_mu[:, 0] + logprob_sigma[:, 0]
joint_prob_tf = tf.exp(log_joint_prod - tf.reduce_max(log_joint_prod))