Constrain tensorflow probability to positive coefficients

Question

I have a tensorflow sts model I wish to constrain the Linear Regression coefficients to greater than zero. I understand this can be achieved by passing a HalfNormal distribution as a prior:

network_effects = tfp.sts.LinearRegression(
  design_matrix=tf.stack((df-df.mean()).values.astype(np.float32)),
  name='network_effects',
  weights_prior=tfd.HalfNormal(scale=2.0))
    
autoregressive = sts.Autoregressive(
  order=8,
  observed_time_series=observed_time_series,
  name='autoregressive')

However, it complains that my dtypes are not the same with the error:

ValueError: SampleHalfNormal, type=<dtype: 'float32'>, must be of the same type (<dtype: 'float64'>) as design_matrix_linop.

Is my method of constraining the Linear Regressor coefficients correct and if so, how do I specify that the HalfNormal distribution is of type float64?

Answer 1

Not at all obvious how this is done, but I'm not overly familiar with Tensorflow. However the answer is to set the scale argument to float64 as follows:

tfd.HalfNormal(scale=np.float64(2.0))

The answer came from this post here: How can I create an array of distributions in TensorFlow Probability?