TensorFlow 2.0 Keras layers with custom tensors as variables

Answer based on the comment below

Given you have:

kernel = createTheKernelVarBasedOnWhatYouWant() #shape (784, 64)
bias = createTheBiasVarBasedOnWhatYouWant() #shape (64,)

Make a simple function copying the code from Dense:

def custom_dense(x):
    inputs, kernel, bias = x

    outputs = K.dot(inputs, kernel)
    outputs = K.bias_add(outputs, bias, data_format='channels_last')
    return outputs

Use the function in a Lambda layer:

layer = Lambda(custom_dense)
hiddens = layer([x, kernel, bias])

Warning: kernel and bias must be produced from a Keras layer, or come from an kernel = Input(tensor=the_kernel_var) and bias = Input(tensor=bias_var)

If the warning above is bad for you, you can always use kernel and bias "from outside", like:

def custom_dense(inputs):
    outputs = K.dot(inputs, kernel) #where kernel is not part of the arguments anymore
    outputs = K.bias_add(outputs, bias, data_format='channels_last')
    return outputs

layer = Lambda(custom_dense)
hiddens = layer(x)

This last option makes it a bit more complicated to save/load models.

Old answer

You should probably use a Keras Dense layer and set its weights in a standard way:

layer = tf.keras.layers.Dense(64, name='the_layer')
layer.set_weights([np.random.rand(784, 64), np.random.rand(64)])

If you need that these weights are not trainable, before compiling the keras model you set:

model.get_layer('the_layer').trainable=False

If you want direct access to the variables as tensors, they are:

kernel = layer.kernel    
bias = layer.bias

There are plenty of other options, but that depends on your exact intention, which is not clear in your question.

Below is a general-purpose solution that works with arbitrary Keras models in TF2.

First, we need to define an auxiliary function canonical_variable_name and a context manager custom_make_variable with the following signatures (see implementation in meta-blocks library).

def canonical_variable_name(variable_name: str, outer_scope: str):
    """Returns the canonical variable name: `outer_scope/.../name`."""
    # ...

@contextlib.contextmanager
def custom_make_variable(
    canonical_custom_variables: Dict[str, tf.Tensor], outer_scope: str
):
    """A context manager that overrides `make_variable` with a custom function.

    When building layers, Keras uses `make_variable` function to create weights
    (kernels and biases for each layer). This function wraps `make_variable` with
    a closure that infers the canonical name of the variable being created (of the
    form `outer_scope/.../var_name`) and looks it up in the `custom_variables` dict
    that maps canonical names to tensors. The function adheres the following logic:

    * If there is a match, it does a few checks (shape, dtype, etc.) and returns
      the found tensor instead of creating a new variable.
    * If there is a match but checks fail, it throws an exception.
    * If there are no matching `custom_variables`, it calls the original
      `make_variable` utility function and returns a newly created variable.
    """
    # ...

Using these functions, we can create arbitrary Keras models with custom tensors used as variables:

import numpy as np
import tensorflow as tf

canonical_custom_variables = {
    "model/dense/kernel": tf.constant(
        np.random.rand(784, 64), name="custom_kernel", dtype=tf.float32),
    "model/dense/bias": tf.constant(
        np.random.rand(64), name="custom_bias", dtype=tf.float32),
}

# Compute hiddens using a dense layer with custom variables.
x = tf.random.normal(shape=(1, 784), name="inputs")
with custom_make_variable(canonical_custom_variables, outer_scope="model"):
    Layer = tf.layers.Dense(64)
    hiddens = Layer(x)

print(Layer.variables)

Not entirely sure I understand your question correctly, but it seems to me that it should be possible to do what you want with a combination of custom layers and keras functional api.

Custom layers allow you to build any layer you want in a way that is compatible with Keras, e.g.:

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],
                                      initializer='normal')

        self.bias = self.add_weight("bias", 
                                    shape=[self.num_outputs,],
                                    initializer='normal')

    def call(self, inputs):
        return tf.matmul(inputs, self.kernel) + self.bias

and the functional api allows you to access the outputs of said layers and re-use them:

inputs = keras.Input(shape=(784,), name='img')
x1 = MyDenseLayer(64, activation='relu')(inputs)
x2 = MyDenseLayer(64, activation='relu')(x1)
outputs = MyDenseLayer(10, activation='softmax')(x2)

model = keras.Model(inputs=inputs, outputs=outputs, name='mnist_model')

Here x1 and x2 can be connected to other subnets.