Multi-level model with nonsense state and county coefficients

I think the issue is with dims= ("county", "state")). Try to reframe the model so that this is dims="county", and yield_estimate becomes

yield_estimate = (
                 a_county[county_idx] +
                )