So I think this logic will work but it still looks a bit weird to my eye, because we're checking _original_treatment_featurizer for non-None-ness when we really care about is that we can use transformer instead. I guess we can't check that transformer is non-None because it also exists for discrete treatments but we don't want to transform in that case?

I think perhaps it would be worth being more explicit and promoting "should we use the treatment featurizer when getting marginal effects" to a first-class concept stored as a flag at the base estimator level because that would make this logic much more clear.

It's also a bit odd because currently transformer is defined as an attribute on the TransformerMixin class, but you're using it here in LinearCateEstimator, which might not be a TransformerMixin instance. So maybe that attribute should also be moved up to be part of the BaseCateEstimator API itself (since you could in principle have TransformerMixins that aren't LinearCateEstimators or vice versa, but we'd now want the attribute in each place).