Update uq_benchmark_2019 code to use tfp.math.psd_kernels.FeatureScaled
Update uq_benchmark_2019 code to use tfp.math.psd_kernels.FeatureScaled

This enables deletion of some custom logic and shape hackery to handle
per-dimension scaling.

In order to pass the existing tests we need to use sqrt(length_scale) because
the existing code was computing kernels like

k(x, y) \propto  e^{-(x - y)^2 / \lambda}

instead of

k(x, y) \propto  e^{-((x - y) / \lambda)^2}

the latter of which is what is computed by the combination of
tfp.math.psd_kernels.FeatureScaled and
tfp.math.psd_kernels.ExponentiatedQuadratic (similarly for other length-scaled
kernels like the various Materns)

(This change was precipitated by an upcoming update to VGP which, among other
things, is more sensitive to being sure that kernel batch_shape computations are
correct; the kernels being produced by this code, by not overriding
_batch_shape logic, were coming up with invalid values when computing the batch
shape.)

PiperOrigin-RevId: 286612666

google-research/google-research,google-research/google-research,google-research/google-research,google-research/google-research,google-research/google-research,google-research/google-research,google-research/google-research,google-research/google-research,google-research/google-research,google-research/google-research,google-research/google-research