View source on GitHub |
A layer that uses tf.einsum
as the backing computation.
tf.keras.layers.EinsumDense(
equation,
output_shape,
activation=None,
bias_axes=None,
kernel_initializer='glorot_uniform',
bias_initializer='zeros',
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
**kwargs
)
This layer can perform einsum calculations of arbitrary dimensionality.
Examples:
Biased dense layer with einsums
This example shows how to instantiate a standard Keras dense layer using
einsum operations. This example is equivalent to
tf.keras.layers.Dense(64, use_bias=True)
.
layer = tf.keras.layers.EinsumDense("ab,bc->ac",
output_shape=64,
bias_axes="c")
input_tensor = tf.keras.Input(shape=[32])
output_tensor = layer(input_tensor)
output_tensor
<... shape=(None, 64) dtype=...>
Applying a dense layer to a sequence
This example shows how to instantiate a layer that applies the same dense
operation to every element in a sequence. Here, the output_shape
has two
values (since there are two non-batch dimensions in the output); the first
dimension in the output_shape
is None
, because the sequence dimension
b
has an unknown shape.
layer = tf.keras.layers.EinsumDense("abc,cd->abd",
output_shape=(None, 64),
bias_axes="d")
input_tensor = tf.keras.Input(shape=[32, 128])
output_tensor = layer(input_tensor)
output_tensor
<... shape=(None, 32, 64) dtype=...>
Applying a dense layer to a sequence using ellipses
This example shows how to instantiate a layer that applies the same dense operation to every element in a sequence, but uses the ellipsis notation instead of specifying the batch and sequence dimensions.
Because we are using ellipsis notation and have specified only one axis, the
output_shape
arg is a single value. When instantiated in this way, the
layer can handle any number of sequence dimensions - including the case
where no sequence dimension exists.
layer = tf.keras.layers.EinsumDense("...x,xy->...y",
output_shape=64,
bias_axes="y")
input_tensor = tf.keras.Input(shape=[32, 128])
output_tensor = layer(input_tensor)
output_tensor
<... shape=(None, 32, 64) dtype=...>