Q: How to implement some trainable parameters in the model of Keras like nn.Parameters() in Pytorch? I just wanna to implement some trainable parameters in my model with Keras. In Pytorch, we can do it by using torch.nn.Parameter() like below:
self.a = nn.Parameter(torch.ones(8))
self.b = nn.Parameter(torch.zeros(16,8))

I think by doing this in pytorch it can add some trainable parameters into the model. And now I wanna to know, how to achieve similar operations in keras?
Any suggestions or advice are welcomed!
THX! :)
p.s. I just write a custom layer in Keras as below:
class Mylayer(Layer):

    def __init__(self,input_dim,output_dim,**kwargs):
        self.input_dim = input_dim
        self.output_dim = output_dim
        super(Mylayer,self).__init__(**kwargs)

    def build(self):

        self.kernel = self.add_weight(name='pi',
                                      shape=(self.input_dim,self.output_dim),
                                      initializer='zeros',
                                      trainable=True)
        self.kernel_2 = self.add_weight(name='mean',
                                        shape=(self.input_dim,self.output_dim),
                                        initializer='ones',
                                        trainable=True)

        super(Mylayer,self).build()

    def call(self,x):
        return x,self.kernel,self.kernel_2

and I wanna to know if I haven't change the tensor which pass through the layer, should I write the function def compute_output_shape() for necessary?

A: You need to create the trainable weights in a custom layer:
class MyLayer(Layer):
    def __init__(self, my_args, **kwargs):
        #do whatever you need with my_args

        super(MyLayer, self).__init__(**kwargs) 

    #you create the weights in build:
    def build(self, input_shape):
        #use the input_shape to infer the necessary shapes for weights
        #use self.whatever_you_registered_in_init to help you, like units, etc. 

        self.kernel = self.add_weight(name='kernel', 
                                  shape=the_shape_you_calculated,
                                  initializer='uniform',
                                  trainable=True)

        #create as many weights as necessary for this layer

        #build the layer - equivalent to self.built=True
        super(MyLayer, self).build(input_shape)

    #create the layer operation here
    def call(self, inputs):
        #do whatever operations are needed
        #example:
        return inputs * self.kernel #make sure the shapes are compatible

    #tell keras about the output shape of your layer
    def compute_output_shape(self, input_shape):
        #calculate the output shape based on the input shape and your layer's rules
        return calculated_output_shape

Now use your layer in the model.

If you are using eager execution on with tensorflow and creating a custom training loop, you can work pretty much the same way you do with PyTorch, and you can create weights outside layers with tf.Variable, passing them as parameters to the gradient calculation methods. 
