Adjust ate inference to get the exact stderr when final stage is linear.

[heimengqi]

Add some context for this implementation:
In order to get the analytical inference of E(theta(x)), we need to know stderr of E(theta(x)), however previously we could only get the conservative upper bound of it. Given the fact that only in linear cate scenario, E(theta(x))=theta(E(x)), and we could get STD(theta(E(x))) by calling prediction_stderr(X.mean(axis=0)), which will return the exact standard error of the mean.

For const_marginal_effect we learn theta(x) alone, then the stderr could be adjusted when final stage is linear in both continues and categorical scenario. For effect, in continues scenario we learn cross product of X and T (as XT), we get the stderr by calling prediction_stderr(XT.mean(axis=0)), in categorical scenario we learn the effect of a specific treatment comparing to baseline, for each sample, the effect could be different given what the treatment is. We couldn't learn theta(XT.mean(axis=0)) unless T1 is a constant or a list of constant. So I only adjust effect_inference only when T1 is a constant or list of constant.

[kbattocchi]

A few comments, mostly minor. But shouldn't there also be something that tests that we get back a non-None mean_stderr for the linear case?

[heimengqi]

A few comments, mostly minor. But shouldn't there also be something that tests that we get back a non-None mean_stderr for the linear case?

Sure. I will add a test around that. I also added some context about this PR to explain why I do it in this way. Please also help confirm the whether logic seems right to you.

[kbattocchi]

LGTM