declare-lab/tango2

1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14import math15 16import flax.linen as nn17import jax.numpy as jnp18 19 20def get_sinusoidal_embeddings(21    timesteps: jnp.ndarray,22    embedding_dim: int,23    freq_shift: float = 1,24    min_timescale: float = 1,25    max_timescale: float = 1.0e4,26    flip_sin_to_cos: bool = False,27    scale: float = 1.0,28) -> jnp.ndarray:29    """Returns the positional encoding (same as Tensor2Tensor).30 31    Args:32        timesteps: a 1-D Tensor of N indices, one per batch element.33        These may be fractional.34        embedding_dim: The number of output channels.35        min_timescale: The smallest time unit (should probably be 0.0).36        max_timescale: The largest time unit.37    Returns:38        a Tensor of timing signals [N, num_channels]39    """40    assert timesteps.ndim == 1, "Timesteps should be a 1d-array"41    assert embedding_dim % 2 == 0, f"Embedding dimension {embedding_dim} should be even"42    num_timescales = float(embedding_dim // 2)43    log_timescale_increment = math.log(max_timescale / min_timescale) / (num_timescales - freq_shift)44    inv_timescales = min_timescale * jnp.exp(jnp.arange(num_timescales, dtype=jnp.float32) * -log_timescale_increment)45    emb = jnp.expand_dims(timesteps, 1) * jnp.expand_dims(inv_timescales, 0)46 47    # scale embeddings48    scaled_time = scale * emb49 50    if flip_sin_to_cos:51        signal = jnp.concatenate([jnp.cos(scaled_time), jnp.sin(scaled_time)], axis=1)52    else:53        signal = jnp.concatenate([jnp.sin(scaled_time), jnp.cos(scaled_time)], axis=1)54    signal = jnp.reshape(signal, [jnp.shape(timesteps)[0], embedding_dim])55    return signal56 57 58class FlaxTimestepEmbedding(nn.Module):59    r"""60    Time step Embedding Module. Learns embeddings for input time steps.61 62    Args:63        time_embed_dim (`int`, *optional*, defaults to `32`):64                Time step embedding dimension65        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):66                Parameters `dtype`67    """68    time_embed_dim: int = 3269    dtype: jnp.dtype = jnp.float3270 71    @nn.compact72    def __call__(self, temb):73        temb = nn.Dense(self.time_embed_dim, dtype=self.dtype, name="linear_1")(temb)74        temb = nn.silu(temb)75        temb = nn.Dense(self.time_embed_dim, dtype=self.dtype, name="linear_2")(temb)76        return temb77 78 79class FlaxTimesteps(nn.Module):80    r"""81    Wrapper Module for sinusoidal Time step Embeddings as described in https://arxiv.org/abs/2006.1123982 83    Args:84        dim (`int`, *optional*, defaults to `32`):85                Time step embedding dimension86    """87    dim: int = 3288    flip_sin_to_cos: bool = False89    freq_shift: float = 190 91    @nn.compact92    def __call__(self, timesteps):93        return get_sinusoidal_embeddings(94            timesteps, embedding_dim=self.dim, flip_sin_to_cos=self.flip_sin_to_cos, freq_shift=self.freq_shift95        )96