How to add a custom layer after a densevariational layer?

I have made a small bayesian neural network with few dense variational layers.

import numpy as np 
from sklearn.model_selection import train_test_split 
from tqdm.notebook import tqdm 
import tensorflow_probability as tfp 
import tensorflow as tf 
from tensorflow.keras.layers import Input 
from tensorflow.keras.optimizers import Adam 
tfd= tfp.distributions 

def prior(kernel_size, bias_size, dtype=None):
    n = kernel_size + bias_size
    # Independent Normal Distribution
    return lambda t: tfd.Independent(tfd.Normal(loc=tf.zeros(n, dtype=dtype),scale=1), reinterpreted_batch_ndims=1)

train_size = X_train.shape[0] batch_size = 256

def NLL(y, distr): 
    return -distr.log_prob(y) 
def normal_sp(params):
    return tfd.Normal(loc=params[:,:1],scale=1e-5 + 0.001*tf.keras.backend.exp(params[:,1:]))# both parameters are learnable


def random_gaussian_initializer(shape, dtype):
    n = int(shape / 2)
    loc_norm = tf.random_normal_initializer(mean=0., stddev=0.1)
    loc = tf.Variable(initial_value=loc_norm(shape=(n,), dtype=dtype) )
    scale_norm = tf.random_normal_initializer(mean=-3., stddev=0.1)
    scale = tf.Variable(initial_value=scale_norm(shape=(n,), dtype=dtype))
    return tf.concat([loc, scale], 0)

def posterior_mean_field(kernel_size, bias_size=0, dtype=None):
    n = kernel_size + bias_size
    return tf.keras.Sequential([
      tfp.layers.VariableLayer(2 * n, dtype=dtype,initializer=lambda shape, dtype: random_gaussian_initializer(shape, dtype), trainable=True),
      tfp.layers.DistributionLambda(lambda t: tfd.Independent(tfd.Normal(loc=t[..., :n],scale=1e-5 + 0.02*tf.nn.softplus(0.04 + t[..., n:])),reinterpreted_batch_ndims=1)),
    ])

def create_probabilistic_bnn_model(train_size):

    model = tf.keras.Sequential([
        tf.keras.Input(shape=(1,)),
        tfp.layers.DenseVariational(50, posterior_mean_field, prior, kl_weight=1/train_size, activation='relu', kl_use_exact=True),
        tfp.layers.DenseVariational(50, posterior_mean_field, prior, kl_weight=1/train_size, activation='relu', kl_use_exact=True),
        tfp.layers.DenseVariational(50, posterior_mean_field, prior, kl_weight=1/train_size, activation='relu', kl_use_exact=True),
        tfp.layers.DenseVariational(2, posterior_mean_field, prior, kl_weight=1/train_size, kl_use_exact=True),
        tfp.layers.DistributionLambda(normal_sp)


    ])
    return model

train_size = x.shape[0]
batch_size = 256

optimizer = tf.optimizers.Adam(learning_rate=0.0002)

keras_BNN = create_probabilistic_bnn_model(train_size=train_size)
keras_BNN.compile(optimizer=optimizer,loss=NLL) 

I taught to add a custom layer at the end so that the bnn would learn not directly the target y but a parameter w of a simple equation $y= f(x,w)$.

  def __init__(self, num_outputs):
    super(MyDenseLayer, self).__init__()
    self.num_outputs = num_outputs

  def build(self, input_shape):
    self.kernel = self.add_weight("kernel",
                                  shape=[int(input_shape[-1]),
                                         self.num_outputs])
    
 

  def call(self, inputs):
    q0= 3/2*inputs[0]
    j0=1
    sol= 0.5*(1-q0)*x -1/24(7-10*q0 -9*q0**2 )*x^2 +j0
 
    return  sol

layer = MyDenseLayer(1)

But when I tried to add MyDenseLayer(1) after tfp.layers.DistributionLambda(normal_sp) I got the error message:

TypeError: Value passed to parameter 'x' has DataType int32 not in list of allowed 
values: bfloat16, float16, float32, float64, complex64, complex128

What do I do wrong?

You should stick to one data type when doing your calculations and it should be fine. Here is an example, but I cannot verify that the logic is correct (that is up to you):

import tensorflow as tf
import tensorflow_probability as tfp

tfk = tf.keras
tfkl = tf.keras.layers
tfd = tfp.distributions
tfpl = tfp.layers

class MyDenseLayer(tf.keras.layers.Layer):
  def __init__(self, num_outputs):
    super(MyDenseLayer, self).__init__()
    self.num_outputs = num_outputs

  def build(self, input_shape):
    self.kernel = self.add_weight("kernel",
                                  shape=[int(input_shape[-1]),
                                         self.num_outputs])
    
 

  def call(self, inputs):
    x = 4
    q0= 3/2*inputs[0]
    j0 = 1
    xor = tf.cast(x, dtype=tf.int32)^2
    sol= 0.5*(1-q0)*x -1/24 * (7-10*q0 -9*q0**2 )* tf.cast(xor, dtype=tf.float32) +j0

    return  sol


def normal_sp(params):
    return tfd.Normal(loc=params[:,:1],scale=1e-5 + 0.001*tf.keras.backend.exp(params[:,1:]))# both parameters are learnable

layer = MyDenseLayer(1)
model = tfk.Sequential([
  tfkl.Dense(2, input_dim=2),
  tfpl.DistributionLambda(normal_sp),
  layer
  ])

model(tf.random.normal((2, 2)))

Note that ^ that does not work with floats that is why I am explicitly casting to int, assuming that x is later some float tensor. If not, you do not have to cast.