@model function λ(y)
lambda1 ~ Exponential(alpha)
lambda2 ~ Exponential(alpha)
τ ~ Uniform(1, n_count_data)
for i in 1:length(y)
y[i] ~ Poisson(τ > i ? lambda1 : lambda2)
end
# return τ
end
e = sample(λ(count_data), HMC(0.05,10), 10000)
savefig(plot(e), "sample.png")

Every time I run the program (on the same data) the plotted sample distributions looks very different.
I have upload an example (but as new user I can embedded only one file):
W

Looks like that step size is too large, try to use HMC(0.01, 10) or something like that, or use NUTS() to let it tune step size automatically for you. A MWE with data is better.

Thank you @yiyuezhuo for the advice.
Now the sample distributions looks more stable (but often their density converges in different points with HMC()) but the result is different from what I expect.
I’m trying to wrote in Julia this PyMC example:

And my posterior distributions looks very different:
(with NUTS() sampler)

pm.Exponential uses a different parameterization compared to Distributions or Turing, so you should specify alpha = mean(count_data) instead of alpha = 1 / mean(count_data).

tau is in fact a discrete value, though you can use a continuous approximation in some sense, the HMC or gradient will not work for tau. You should specify a so-called CompoundStep in terms of pymc.

For example, if you replace the pm.Metropolis with pm.HamiltonianMC, you will find that parameters use different sampler:

### Mysterious code to be explained in Chapter 3.
with model:
step = pm.Metropolis()
trace = pm.sample(10000, tune=5000,step=step)
"""
Sequential sampling (2 chains in 1 job)
CompoundStep
>Metropolis: [tau]
>Metropolis: [lambda_2]
>Metropolis: [lambda_1]
100%|██████████| 15000/15000 [00:06<00:00, 2194.08it/s]
100%|██████████| 15000/15000 [00:06<00:00, 2469.70it/s]
The number of effective samples is smaller than 25% for some parameters.
"""
### Mysterious code to be explained in Chapter 3.
with model:
step = pm.HamiltonianMC()
trace = pm.sample(10000, tune=5000,step=step)
"""
Sequential sampling (2 chains in 1 job)
CompoundStep
>HamiltonianMC: [lambda_2, lambda_1]
>Metropolis: [tau]
100%|██████████| 15000/15000 [00:07<00:00, 1899.67it/s]
100%|██████████| 15000/15000 [00:07<00:00, 2078.96it/s]
The estimated number of effective samples is smaller than 200 for some parameters.
"""

My solution, while I don’t know much how Gibbs work and MH is not exactly equivalent to pymc.Metropolis:

# alpha = 1 / mean(count_data)
alpha = mean(count_data)
@model function λ(y)
lambda1 ~ Exponential(alpha)
lambda2 ~ Exponential(alpha)
τ ~ Uniform(1, n_count_data)
for i in 1:length(y)
y[i] ~ Poisson(τ > i ? lambda1 : lambda2)
end
# return τ
end
sampler = Gibbs(MH(:τ), HMC(0.01, 5, :lambda1, :lambda2))
e = sample(λ(count_data), sampler, 10000)
plot(e)

With alpha=mean(count_data) result are better.
With Gibbs() sampler I obtain always valid samples.
Using HMC(0.01, 5) sometimes I obtain something like this: