How to use a Tensorflow (python) custom layer with only one trainable variable?

Question

I am looking for a way to add a custom layer to my TensorFlow (python) model. This layer should consist of one trainable variable delta. This variable has to be used as an argument of the bessel function special.jn(n, delta). Let input_vector be the d-dimensional vector containing my inputs. I want this to be multiplied to bessel_tensor. The latter should be an (d x d)-Tensor containing the bessel functions for different orders n but all with the same trainable delta. Now, I want to train for that delta to be optimized.

How can I achieve that?

What I've done so far:

• I created an custom layer with the desired trainable weight delta
• I tried to create the Tensor(/Matrix) described above. This did not work because I've not been able to adress the cells individually.
• Using the funciton special.jn(n, delta) in the context of tensors did not work quite well so I created a sum approaching the desired value

A minimal reproducible example would be:

import tensorflow as tf 
from tensorflow.keras import layers
from scipy import special
import pandas as pd
import math as math

# Define parameters
epochs = 12
batch_size = 100
test_size = 1000
training_size = 3000

# Loading training data
test_data = pd.read_csv("data/test_data.csv", delimiter=";")
test_labels = pd.read_csv("data/test_labels.csv")

train_data = pd.read_csv("data/train_data.csv", delimiter=";")
train_labels = pd.read_csv("data/train_labels.csv")

test_data = test_data[:test_size]
test_labels = test_labels[:test_size]

train_data = train_data[:training_size]
train_labels = train_labels[:training_size]


# Layer class
class BesselLayer(layers.Layer):
    def __init__(self, **kwargs):
        super(BesselLayer, self).__init__(**kwargs)


    def build(self, input_shape):
        self.delta = self.add_weight(name="delta", shape=(1, 1), initializer="random_normal", trainable=True)


    def call(self, input_vector):
        
        # Should be a tensor of shape (d x d) with d = 2 shown below
        bessel_tensor = tf.Tensor([[special.jn(0, self.delta), special.jn(1, self.delta)], [special.jn(2, self.delta), special.jn(3, self.delta)]])

        return tf.matmul(input_vector, bessel_tensor)




# Creation of the model
model = tf.keras.models.Sequential()
model.add(tf.keras.layers.Dense(784, activation=tf.nn.relu, input_shape=(784,))) # Dense layer 784
bessel_layer = BesselLayer()
model.add(bessel_layer) # Custom defined layer 784
model.add(tf.keras.layers.Dense(10, activation=tf.nn.softmax)) # Output layer representing the output data


# Compile the model
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=['accuracy'])
print(model.summary())


# Fit the model
history = model.fit(train_data, train_labels, epochs=epochs, batch_size=batch_size)

# Evaluate the model
test_loss, test_accuracy = model.evaluate(test_data, test_labels)

This does not work because a trainable Tensor has to be defined somehow different