A Simple 3-Variable Toy Problem

I’m just getting up to speed with Pyro and probabilistic modeling, so bear with the (perhaps) basic question.
A problem I’m working on is analogous to the following game.

Let’s say there is a bag of coins, some red and some blue. From your sample you can estimate the frequency with which red coins appear using the total number of coins and the number of red coins therein.

The red coins do not all have the same chance of flipping heads, but the parameter for each is drawn from a common distribution (let’s say a beta distribution). Likewise for the blue – they each have a different chance of heads, but all drawn from the same distribution.

You have a sample set of n coins, and each coin is flipped a different number of times. The color of the coin, number of successes, and number of trials is recorded for each coin.

Now you draw a new coin from the bag, but do not know the color. It is flipped some number of times and produces some number of successes. You want to estimate the probability that the coin is red given that result.

How should this problem be modeled?

Here are my thoughts – but I’m a bit unclear of their correctness or of the right implementation as a model in Pyro.

  • Would the prior for the color of the coin be a binomial distribution, or dirichlet distribution?
  • An individual coin’s chance of heads would be a binomial parameter drawn from the beta distribution for the color, right? So there would be two beta distributions indexed by the color.
  • An individual result of successes and total trials for a coin would come from a binomial distribution with the probability of success above.

Am I on the right track?