human
I am working on the file "Annealed importance sampling.ipynb".
The first task is:
2. Define annealing distributions
Here we'll be annealing between two unnormalized Gaussian distributions with geometric mean intermediates.
They have the same variances, so they should have the same normalizing constants.
Can you write Python program for it?

=====
Give me a the Python program you requested?
assistant
#Here is the Python program you requested.

num_intermediates = 25
betas = np.linspace(0,1,num_intermediates+2)
dim=1

def initial_density(x):
    return np.exp(-((x)**2).sum()/2)

def draw_from_initial():
    return npr.randn(dim)

def target_density(x):
    return np.exp(-((x-4)**2).sum()/2)
    
class GeometricMean():
    def __init__(self,initial,target,beta):
        self.initial = initial
        self.target = target
        self.beta = beta
    
    def __call__(self,x):
        f1_x = self.initial(x)
        fT_x = self.target(x)
        return f1_x**(1-self.beta) * fT_x**self.beta

annealing_distributions = [GeometricMean(initial_density,target_density,beta) for beta in betas]