self._intermediate_dense.append(
          tf_keras.layers.EinsumDense(
              "abc,cd->abd",
              output_shape=(None, self._inner_dim),
              bias_axes="d",