"""A UNet model implementation for use with diffusion models

Adapted from OpenAI guided diffusion, with slight modifications

and additional features

https://github.com/openai/guided-diffusion

MIT License

Copyright (c) 2021 OpenAI

Permission is hereby granted, free of charge, to any person obtaining a copy

of this software and associated documentation files (the "Software"), to deal

in the Software without restriction, including without limitation the rights

to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

copies of the Software, and to permit persons to whom the Software is

furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all

copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE

SOFTWARE.

Authors

* Artem Ploujnikov 2022

"""

import math

from abc import abstractmethod

from typing import Optional

import torch

import torch.nn as nn

import torch.nn.functional as F

from speechbrain.utils.data_utils import pad_divisible

from .autoencoders import NormalizingAutoencoder

def fixup(module, use_fixup_init=True):

"""

Zero out the parameters of a module and return it.

Arguments

---------

module: torch.nn.Module

a module

use_fixup_init: bool

whether to zero out the parameters. If set to

false, the function is a no-op

Returns

-------

The fixed module

"""

if use_fixup_init:

for p in module.parameters():

p.detach().zero_()

return module

def conv_nd(dims, *args, **kwargs):

"""

Create a 1D, 2D, or 3D convolution module.

Arguments

---------

dims: int

The number of dimensions

*args: tuple

**kwargs: dict

Any remaining arguments are passed to the constructor

Returns

-------

The constructed Conv layer

"""

if dims == 1:

return nn.Conv1d(*args, **kwargs)

elif dims == 2:

return nn.Conv2d(*args, **kwargs)

elif dims == 3:

return nn.Conv3d(*args, **kwargs)

raise ValueError(f"unsupported dimensions: {dims}")

def avg_pool_nd(dims, *args, **kwargs):

"""

Create a 1D, 2D, or 3D average pooling module.

"""

if dims == 1:

return nn.AvgPool1d(*args, **kwargs)

elif dims == 2:

return nn.AvgPool2d(*args, **kwargs)

elif dims == 3:

return nn.AvgPool3d(*args, **kwargs)

raise ValueError(f"unsupported dimensions: {dims}")

def timestep_embedding(timesteps, dim, max_period=10000):

"""

Create sinusoidal timestep embeddings.

Arguments

---------

timesteps: torch.Tensor

a 1-D Tensor of N indices, one per batch element. These may be fractional.

dim: int

the dimension of the output.

max_period: int

controls the minimum frequency of the embeddings.

Returns

-------

result: torch.Tensor

an [N x dim] Tensor of positional embeddings.

"""

half = dim // 2

freqs = torch.exp(

-math.log(max_period)

* torch.arange(start=0, end=half, dtype=torch.float32)

/ half

).to(device=timesteps.device)

args = timesteps[:, None].float() * freqs[None]

embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)

if dim % 2:

embedding = torch.cat(

[embedding, torch.zeros_like(embedding[:, :1])], dim=-1

)

return embedding

class AttentionPool2d(nn.Module):

"""Two-dimensional attentional pooling

Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py

Arguments

---------

spatial_dim: int

the size of the spatial dimension

embed_dim: int

the embedding dimension

num_heads_channels: int

the number of attention heads

output_dim: int

the output dimension

Example

-------

>>> attn_pool = AttentionPool2d(

... spatial_dim=64, embed_dim=16, num_heads_channels=2, output_dim=4

... )

>>> x = torch.randn(4, 1, 64, 64)

>>> x_pool = attn_pool(x)

>>> x_pool.shape

torch.Size([4, 4])

"""

def __init__(

self,

spatial_dim: int,

embed_dim: int,

num_heads_channels: int,

output_dim: Optional[int] = None,

):

super().__init__()

self.positional_embedding = nn.Parameter(

torch.randn(embed_dim, spatial_dim**2 + 1) / embed_dim**0.5

)

self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)

self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)

self.num_heads = embed_dim // num_heads_channels

self.attention = QKVAttention(self.num_heads)

def forward(self, x):

"""Computes the attention forward pass

Arguments

---------

x: torch.Tensor

the tensor to be attended to

Returns

-------

result: torch.Tensor

the attention output

"""

b, c, *_spatial = x.shape

x = x.reshape(b, c, -1) # NC(HW)

x = torch.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1)

x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1)

x = self.qkv_proj(x)

x = self.attention(x)

x = self.c_proj(x)

return x[:, :, 0]

class TimestepBlock(nn.Module):

"""

Any module where forward() takes timestep embeddings as a second argument.

"""

@abstractmethod

def forward(self, x, emb=None):

"""

Apply the module to `x` given `emb` timestep embeddings.

Arguments

---------

x: torch.Tensor

the data tensor

emb: torch.Tensor

the embedding tensor

"""

class TimestepEmbedSequential(nn.Sequential, TimestepBlock):

"""A sequential module that passes timestep embeddings to the children that

support it as an extra input.

Example

-------

>>> from speechbrain.nnet.linear import Linear

>>> class MyBlock(TimestepBlock):

... def __init__(self, input_size, output_size, emb_size):

... super().__init__()

... self.lin = Linear(n_neurons=output_size, input_size=input_size)

... self.emb_proj = Linear(

... n_neurons=output_size,

... input_size=emb_size,

... )

...

... def forward(self, x, emb):

... return self.lin(x) + self.emb_proj(emb)

>>> tes = TimestepEmbedSequential(

... MyBlock(128, 64, 16), Linear(n_neurons=32, input_size=64)

... )

>>> x = torch.randn(4, 10, 128)

>>> emb = torch.randn(4, 10, 16)

>>> out = tes(x, emb)

>>> out.shape

torch.Size([4, 10, 32])

"""

def forward(self, x, emb=None):

"""Computes a sequential pass with sequential embeddings where applicable

Arguments

---------

x: torch.Tensor

the data tensor

emb: torch.Tensor

timestep embeddings

Returns

-------

The processed input

"""

for layer in self:

if isinstance(layer, TimestepBlock):

x = layer(x, emb)

else:

x = layer(x)

return x

class Upsample(nn.Module):

"""

An upsampling layer with an optional convolution.

Arguments

---------

channels: torch.Tensor

channels in the inputs and outputs.

use_conv: bool

a bool determining if a convolution is applied.

dims: int

determines if the signal is 1D, 2D, or 3D. If 3D, then

upsampling occurs in the inner-two dimensions.

out_channels: int

Number of output channels. If None, same as input channels.

Example

-------

>>> ups = Upsample(channels=4, use_conv=True, dims=2, out_channels=8)

>>> x = torch.randn(8, 4, 32, 32)

>>> x_up = ups(x)

>>> x_up.shape

torch.Size([8, 8, 64, 64])

"""

def __init__(self, channels, use_conv, dims=2, out_channels=None):

super().__init__()

self.channels = channels

self.out_channels = out_channels or channels

self.use_conv = use_conv

self.dims = dims

if use_conv:

self.conv = conv_nd(

dims, self.channels, self.out_channels, 3, padding=1

)

def forward(self, x):

"""Computes the upsampling pass

Arguments

---------

x: torch.Tensor

layer inputs

Returns

-------

result: torch.Tensor

upsampled outputs"""

assert x.shape[1] == self.channels

if self.dims == 3:

x = F.interpolate(

x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"

)

else:

x = F.interpolate(x, scale_factor=2, mode="nearest")

if self.use_conv:

x = self.conv(x)

return x

class Downsample(nn.Module):

"""

A downsampling layer with an optional convolution.

Arguments

---------

channels: int

channels in the inputs and outputs.

use_conv: bool

a bool determining if a convolution is applied.

dims: int

determines if the signal is 1D, 2D, or 3D. If 3D, then

downsampling occurs in the inner-two dimensions.

out_channels: int

Number of output channels. If None, same as input channels.

Example

-------

>>> ups = Downsample(channels=4, use_conv=True, dims=2, out_channels=8)

>>> x = torch.randn(8, 4, 32, 32)

>>> x_up = ups(x)

>>> x_up.shape

torch.Size([8, 8, 16, 16])

"""

def __init__(self, channels, use_conv, dims=2, out_channels=None):

super().__init__()

self.channels = channels

self.out_channels = out_channels or channels

self.use_conv = use_conv

self.dims = dims

stride = 2 if dims != 3 else (1, 2, 2)

if use_conv:

self.op = conv_nd(

dims,

self.channels,

self.out_channels,

3,

stride=stride,

padding=1,

)

else:

assert self.channels == self.out_channels

self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)

def forward(self, x):

"""Computes the downsampling pass

Arguments

---------

x: torch.Tensor

layer inputs

Returns

-------

result: torch.Tensor

downsampled outputs

"""

assert x.shape[1] == self.channels

return self.op(x)

class ResBlock(TimestepBlock):

"""

A residual block that can optionally change the number of channels.

Arguments

---------

channels: int

the number of input channels.

emb_channels: int

the number of timestep embedding channels.

dropout: float

the rate of dropout.

out_channels: int

if specified, the number of out channels.

use_conv: bool

if True and out_channels is specified, use a spatial

convolution instead of a smaller 1x1 convolution to change the

channels in the skip connection.

dims: int

determines if the signal is 1D, 2D, or 3D.

up: bool

if True, use this block for upsampling.

down: bool

if True, use this block for downsampling.

norm_num_groups: int

the number of groups for group normalization

use_fixup_init: bool

whether to use FixUp initialization

Example

-------

>>> res = ResBlock(

... channels=4,

... emb_channels=8,

... dropout=0.1,

... norm_num_groups=2,

... use_conv=True,

... )

>>> x = torch.randn(2, 4, 32, 32)

>>> emb = torch.randn(2, 8)

>>> res_out = res(x, emb)

>>> res_out.shape

torch.Size([2, 4, 32, 32])

"""

def __init__(

self,

channels,

emb_channels,

dropout,

out_channels=None,

use_conv=False,

dims=2,

up=False,

down=False,

norm_num_groups=32,

use_fixup_init=True,

):

super().__init__()

self.channels = channels

self.emb_channels = emb_channels

self.dropout = dropout

self.out_channels = out_channels or channels

self.use_conv = use_conv

self.in_layers = nn.Sequential(

nn.GroupNorm(norm_num_groups, channels),

nn.SiLU(),

conv_nd(dims, channels, self.out_channels, 3, padding=1),

)

self.updown = up or down

if up:

self.h_upd = Upsample(channels, False, dims)

self.x_upd = Upsample(channels, False, dims)

elif down:

self.h_upd = Downsample(channels, False, dims)

self.x_upd = Downsample(channels, False, dims)

else:

self.h_upd = self.x_upd = nn.Identity()

if emb_channels is not None:

self.emb_layers = nn.Sequential(

nn.SiLU(),

nn.Linear(

emb_channels,

self.out_channels,

),

)

else:

self.emb_layers = None

self.out_layers = nn.Sequential(

nn.GroupNorm(norm_num_groups, self.out_channels),

nn.SiLU(),

nn.Dropout(p=dropout),

fixup(

conv_nd(

dims, self.out_channels, self.out_channels, 3, padding=1

),

use_fixup_init=use_fixup_init,

),

)

if self.out_channels == channels:

self.skip_connection = nn.Identity()

elif use_conv:

self.skip_connection = conv_nd(

dims, channels, self.out_channels, 3, padding=1

)

else:

self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)

def forward(self, x, emb=None):

"""

Apply the block to a torch.Tensor, conditioned on a timestep embedding.

Arguments

---------

x: torch.Tensor

an [N x C x ...] Tensor of features.

emb: torch.Tensor

an [N x emb_channels] Tensor of timestep embeddings.

Returns

-------

result: torch.Tensor

an [N x C x ...] Tensor of outputs.

"""

if self.updown:

in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]

h = in_rest(x)

h = self.h_upd(h)

x = self.x_upd(x)

h = in_conv(h)

else:

h = self.in_layers(x)

if emb is not None:

emb_out = self.emb_layers(emb).type(h.dtype)

while len(emb_out.shape) < len(h.shape):

emb_out = emb_out[..., None]

else:

emb_out = torch.zeros_like(h)

h = h + emb_out

h = self.out_layers(h)

return self.skip_connection(x) + h

class AttentionBlock(nn.Module):

"""

An attention block that allows spatial positions to attend to each other.

Originally ported from here, but adapted to the N-d case.

https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.

Arguments

---------

channels: int

the number of channels

num_heads: int

the number of attention heads

num_head_channels: int

the number of channels in each attention head

norm_num_groups: int

the number of groups used for group normalization

use_fixup_init: bool

whether to use FixUp initialization

Example

-------

>>> attn = AttentionBlock(

... channels=8, num_heads=4, num_head_channels=4, norm_num_groups=2

... )

>>> x = torch.randn(4, 8, 16, 16)

>>> out = attn(x)

>>> out.shape

torch.Size([4, 8, 16, 16])

"""

def __init__(

self,

channels,

num_heads=1,

num_head_channels=-1,

norm_num_groups=32,

use_fixup_init=True,

):

super().__init__()

self.channels = channels

if num_head_channels == -1:

self.num_heads = num_heads

else:

assert channels % num_head_channels == 0, (

f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"

)

self.num_heads = channels // num_head_channels

self.norm = nn.GroupNorm(norm_num_groups, channels)

self.qkv = conv_nd(1, channels, channels * 3, 1)

self.attention = QKVAttention(self.num_heads)

self.proj_out = fixup(conv_nd(1, channels, channels, 1), use_fixup_init)

def forward(self, x):

"""Completes the forward pass

Arguments

---------

x: torch.Tensor

the data to be attended to

Returns

-------

result: torch.Tensor

The data, with attention applied

"""

b, c, *spatial = x.shape

x = x.reshape(b, c, -1)

qkv = self.qkv(self.norm(x))

h = self.attention(qkv)

h = self.proj_out(h)

return (x + h).reshape(b, c, *spatial)

class QKVAttention(nn.Module):

"""

A module which performs QKV attention and splits in a different order.

Arguments

---------

n_heads : int

Number of attention heads.

Example

-------

>>> attn = QKVAttention(4)

>>> n = 4

>>> c = 8

>>> h = 64

>>> w = 16

>>> qkv = torch.randn(4, (3 * h * c), w)

>>> out = attn(qkv)

>>> out.shape

torch.Size([4, 512, 16])

"""

def __init__(self, n_heads):

super().__init__()

self.n_heads = n_heads

def forward(self, qkv):

"""Apply QKV attention.

Arguments

---------

qkv: torch.Tensor

an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.

Returns

-------

result: torch.Tensor

an [N x (H * C) x T] tensor after attention.

"""

bs, width, length = qkv.shape

assert width % (3 * self.n_heads) == 0

ch = width // (3 * self.n_heads)

q, k, v = qkv.chunk(3, dim=1)

scale = 1 / math.sqrt(math.sqrt(ch))

weight = torch.einsum(

"bct,bcs->bts",

(q * scale).view(bs * self.n_heads, ch, length),

(k * scale).view(bs * self.n_heads, ch, length),

) # More stable with f16 than dividing afterwards

weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)

a = torch.einsum(

"bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length)

)

return a.reshape(bs, -1, length)

def build_emb_proj(emb_config, proj_dim=None, use_emb=None):

"""Builds a dictionary of embedding modules for embedding

projections

Arguments

---------

emb_config: dict

a configuration dictionary

proj_dim: int

the target projection dimension

use_emb: dict

an optional dictionary of "switches" to turn

embeddings on and off

Returns

-------

result: torch.nn.ModuleDict

a ModuleDict with a module for each embedding

"""

emb_proj = {}

if emb_config is not None:

for key, item_config in emb_config.items():

if use_emb is None or use_emb.get(key):

if "emb_proj" in item_config:

emb_proj[key] = emb_proj

else:

emb_proj[key] = EmbeddingProjection(

emb_dim=item_config["emb_dim"], proj_dim=proj_dim

)

return nn.ModuleDict(emb_proj)

class UNetModel(nn.Module):

"""

The full UNet model with attention and timestep embedding.

Arguments

---------

in_channels: int

channels in the input torch.Tensor.

model_channels: int

base channel count for the model.

out_channels: int

channels in the output torch.Tensor.

num_res_blocks: int

number of residual blocks per downsample.

attention_resolutions: int

a collection of downsample rates at which

attention will take place. May be a set, list, or tuple.

For example, if this contains 4, then at 4x downsampling, attention

will be used.

dropout: float

the dropout probability.

channel_mult: int

channel multiplier for each level of the UNet.

conv_resample: bool

if True, use learned convolutions for upsampling and

downsampling

dims: int

determines if the signal is 1D, 2D, or 3D.

emb_dim: int

time embedding dimension (defaults to model_channels * 4)

cond_emb: dict

embeddings on which the model will be conditioned

Example:

{

"speaker": {

"emb_dim": 256

},

"label": {

"emb_dim": 12

}

}

use_cond_emb: dict

a dictionary with keys corresponding to keys in cond_emb

and values corresponding to Booleans that turn embeddings

on and off. This is useful in combination with hparams files

to turn embeddings on and off with simple switches

Example:

{"speaker": False, "label": True}

num_heads: int

the number of attention heads in each attention layer.

num_head_channels: int

if specified, ignore num_heads and instead use

a fixed channel width per attention head.

num_heads_upsample: int

works with num_heads to set a different number

of heads for upsampling. Deprecated.

norm_num_groups: int

Number of groups in the norm, default 32

resblock_updown: bool

use residual blocks for up/downsampling.

use_fixup_init: bool

whether to use FixUp initialization

Example

-------

>>> model = UNetModel(

... in_channels=3,

... model_channels=32,

... out_channels=1,

... num_res_blocks=1,

... attention_resolutions=[1],

... )

>>> x = torch.randn(4, 3, 16, 32)

>>> ts = torch.tensor([10, 100, 50, 25])

>>> out = model(x, ts)

>>> out.shape

torch.Size([4, 1, 16, 32])

"""

def __init__(

self,

in_channels,

model_channels,

out_channels,

num_res_blocks,

attention_resolutions,

dropout=0,

channel_mult=(1, 2, 4, 8),

conv_resample=True,

dims=2,

emb_dim=None,

cond_emb=None,

use_cond_emb=None,

num_heads=1,

num_head_channels=-1,

num_heads_upsample=-1,

norm_num_groups=32,

resblock_updown=False,

use_fixup_init=True,

):

super().__init__()

if num_heads_upsample == -1:

num_heads_upsample = num_heads

self.in_channels = in_channels

self.model_channels = model_channels

self.out_channels = out_channels

self.num_res_blocks = num_res_blocks

self.attention_resolutions = attention_resolutions

self.dropout = dropout

self.channel_mult = channel_mult

self.conv_resample = conv_resample

self.dtype = torch.float32

self.num_heads = num_heads

self.num_head_channels = num_head_channels

self.num_heads_upsample = num_heads_upsample

self.cond_emb = cond_emb

self.use_cond_emb = use_cond_emb

if emb_dim is None:

emb_dim = model_channels * 4

self.time_embed = EmbeddingProjection(model_channels, emb_dim)

self.cond_emb_proj = build_emb_proj(

emb_config=cond_emb, proj_dim=emb_dim, use_emb=use_cond_emb

)

ch = input_ch = int(channel_mult[0] * model_channels)

self.input_blocks = nn.ModuleList(

[

TimestepEmbedSequential(

conv_nd(dims, in_channels, ch, 3, padding=1)

)

]

)

self._feature_size = ch

input_block_chans = [ch]

ds = 1

for level, mult in enumerate(channel_mult):

for _ in range(num_res_blocks):

layers = [

ResBlock(

ch,

emb_dim,

dropout,

out_channels=int(mult * model_channels),

dims=dims,

norm_num_groups=norm_num_groups,

use_fixup_init=use_fixup_init,

)

]

ch = int(mult * model_channels)

if ds in attention_resolutions:

layers.append(

AttentionBlock(

ch,

num_heads=num_heads,

num_head_channels=num_head_channels,

norm_num_groups=norm_num_groups,

use_fixup_init=use_fixup_init,

)

)

self.input_blocks.append(TimestepEmbedSequential(*layers))

self._feature_size += ch

input_block_chans.append(ch)

if level != len(channel_mult) - 1:

out_ch = ch

self.input_blocks.append(

TimestepEmbedSequential(

ResBlock(

ch,

emb_dim,

dropout,

out_channels=out_ch,

dims=dims,

down=True,

norm_num_groups=norm_num_groups,

use_fixup_init=use_fixup_init,

)

if resblock_updown

else Downsample(

ch, conv_resample, dims=dims, out_channels=out_ch

)

)

)

ch = out_ch

input_block_chans.append(ch)

ds *= 2

self._feature_size += ch

self.middle_block = TimestepEmbedSequential(

ResBlock(

ch,

emb_dim,

dropout,

dims=dims,

norm_num_groups=norm_num_groups,

use_fixup_init=use_fixup_init,

),

AttentionBlock(

ch,

num_heads=num_heads,

num_head_channels=num_head_channels,

norm_num_groups=norm_num_groups,

use_fixup_init=use_fixup_init,

),

ResBlock(

ch,

emb_dim,

dropout,

dims=dims,

norm_num_groups=norm_num_groups,

use_fixup_init=use_fixup_init,

),

)

self._feature_size += ch

self.output_blocks = nn.ModuleList([])

for level, mult in list(enumerate(channel_mult))[::-1]:

for i in range(num_res_blocks + 1):

ich = input_block_chans.pop()

layers = [

ResBlock(

ch + ich,

emb_dim,

dropout,

out_channels=int(model_channels * mult),

dims=dims,

norm_num_groups=norm_num_groups,

use_fixup_init=use_fixup_init,

)

]

ch = int(model_channels * mult)

if ds in attention_resolutions:

layers.append(

AttentionBlock(

ch,

num_heads=num_heads_upsample,

num_head_channels=num_head_channels,

norm_num_groups=norm_num_groups,

use_fixup_init=use_fixup_init,

)

)

if level and i == num_res_blocks:

out_ch = ch

layers.append(

ResBlock(

ch,

emb_dim,

dropout,

out_channels=out_ch,

dims=dims,

up=True,

norm_num_groups=norm_num_groups,

use_fixup_init=use_fixup_init,

)

if resblock_updown

else Upsample(

ch, conv_resample, dims=dims, out_channels=out_ch

)

)

ds //= 2

self.output_blocks.append(TimestepEmbedSequential(*layers))

self._feature_size += ch

self.out = nn.Sequential(

nn.GroupNorm(norm_num_groups, ch),

nn.SiLU(),

fixup(

conv_nd(dims, input_ch, out_channels, 3, padding=1),

use_fixup_init=use_fixup_init,

),

)

def forward(self, x, timesteps, cond_emb=None):

"""Apply the model to an input batch.

Arguments

---------

x: torch.Tensor

an [N x C x ...] Tensor of inputs.

timesteps: torch.Tensor

a 1-D batch of timesteps.

cond_emb: dict

a string -> tensor dictionary of conditional

embeddings (multiple embeddings are supported)

Returns

-------

result: torch.Tensor