`Estimator`s make too many assumptions about the shapes of outputs.

[josephdviviano]

Right now, the estimator output dim checker only looks at the final output dimension. It is implicitly assuming we are calculating something that looks like a 1-d distribution. I believe this might be too restricting for all possible Estimator use cases, but it is not currently a priority.

[saleml]

Thanks for raising this. Have you ever stumbled upon a case where an error was raised for this reason?