prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  [Flax] 🚨 0.7.0 not working 🚨
  ### Describe the bug
  
  ![Screenshot 2022-11-05 083749-x](https://user-images.githubusercontent.com/54370274/200104662-894bbbce-b435-48fc-af87-b539595fc898.jpg)
  
  
  ### Reproduction
  
  _No response_
  
  ### Logs
  
  _No response_
  
  ### System Info
  
  TPU v3-8
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `src/diffusers/models/embeddings_flax.py`
  Content:
  ```
  1 # Copyright 2022 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 at
  6 #
  7 #     http://www.apache.org/licenses/LICENSE-2.0
  8 #
  9 # Unless required by applicable law or agreed to in writing, software
  10 # 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 and
  13 # limitations under the License.
  14 import math
  15 
  16 import flax.linen as nn
  17 import jax.numpy as jnp
  18 
  19 
  20 def 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     Args:
  31         timesteps: a 1-D Tensor of N indices, one per batch element.
  32         These may be fractional.
  33         embedding_dim: The number of output channels.
  34         min_timescale: The smallest time unit (should probably be 0.0).
  35         max_timescale: The largest time unit.
  36     Returns:
  37         a Tensor of timing signals [N, num_channels]
  38     """
  39     assert timesteps.ndim == 1, "Timesteps should be a 1d-array"
  40     assert embedding_dim % 2 == 0, f"Embedding dimension {embedding_dim} should be even"
  41     num_timescales = float(embedding_dim // 2)
  42     log_timescale_increment = math.log(max_timescale / min_timescale) / (num_timescales - freq_shift)
  43     inv_timescales = min_timescale * jnp.exp(jnp.arange(num_timescales, dtype=jnp.float32) * -log_timescale_increment)
  44     emb = jnp.expand_dims(timesteps, 1) * jnp.expand_dims(inv_timescales, 0)
  45 
  46     # scale embeddings
  47     scaled_time = scale * emb
  48 
  49     if flip_sin_to_cos:
  50         signal = jnp.concatenate([jnp.cos(scaled_time), jnp.sin(scaled_time)], axis=1)
  51     else:
  52         signal = jnp.concatenate([jnp.sin(scaled_time), jnp.cos(scaled_time)], axis=1)
  53     signal = jnp.reshape(signal, [jnp.shape(timesteps)[0], embedding_dim])
  54     return signal
  55 
  56 
  57 class FlaxTimestepEmbedding(nn.Module):
  58     r"""
  59     Time step Embedding Module. Learns embeddings for input time steps.
  60 
  61     Args:
  62         time_embed_dim (`int`, *optional*, defaults to `32`):
  63                 Time step embedding dimension
  64         dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):
  65                 Parameters `dtype`
  66     """
  67     time_embed_dim: int = 32
  68     dtype: jnp.dtype = jnp.float32
  69 
  70     @nn.compact
  71     def __call__(self, temb):
  72         temb = nn.Dense(self.time_embed_dim, dtype=self.dtype, name="linear_1")(temb)
  73         temb = nn.silu(temb)
  74         temb = nn.Dense(self.time_embed_dim, dtype=self.dtype, name="linear_2")(temb)
  75         return temb
  76 
  77 
  78 class FlaxTimesteps(nn.Module):
  79     r"""
  80     Wrapper Module for sinusoidal Time step Embeddings as described in https://arxiv.org/abs/2006.11239
  81 
  82     Args:
  83         dim (`int`, *optional*, defaults to `32`):
  84                 Time step embedding dimension
  85     """
  86     dim: int = 32
  87     freq_shift: float = 1
  88 
  89     @nn.compact
  90     def __call__(self, timesteps):
  91         return get_sinusoidal_embeddings(timesteps, embedding_dim=self.dim, freq_shift=self.freq_shift)
  92
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/src/diffusers/models/embeddings_flax.py b/src/diffusers/models/embeddings_flax.py
  --- a/src/diffusers/models/embeddings_flax.py
  +++ b/src/diffusers/models/embeddings_flax.py
  @@ -88,4 +88,6 @@
   
       @nn.compact
       def __call__(self, timesteps):
  -        return get_sinusoidal_embeddings(timesteps, embedding_dim=self.dim, freq_shift=self.freq_shift)
  +        return get_sinusoidal_embeddings(
  +            timesteps, embedding_dim=self.dim, freq_shift=self.freq_shift, flip_sin_to_cos=True
  +        )
  
verification_info: {"golden_diff": "diff --git a/src/diffusers/models/embeddings_flax.py b/src/diffusers/models/embeddings_flax.py\n--- a/src/diffusers/models/embeddings_flax.py\n+++ b/src/diffusers/models/embeddings_flax.py\n@@ -88,4 +88,6 @@\n \n     @nn.compact\n     def __call__(self, timesteps):\n-        return get_sinusoidal_embeddings(timesteps, embedding_dim=self.dim, freq_shift=self.freq_shift)\n+        return get_sinusoidal_embeddings(\n+            timesteps, embedding_dim=self.dim, freq_shift=self.freq_shift, flip_sin_to_cos=True\n+        )\n", "issue": "[Flax] \ud83d\udea8 0.7.0 not working \ud83d\udea8\n### Describe the bug\n\n![Screenshot 2022-11-05 083749-x](https://user-images.githubusercontent.com/54370274/200104662-894bbbce-b435-48fc-af87-b539595fc898.jpg)\r\n\n\n### Reproduction\n\n_No response_\n\n### Logs\n\n_No response_\n\n### System Info\n\nTPU v3-8\n", "before_files": [{"content": "# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport math\n\nimport flax.linen as nn\nimport jax.numpy as jnp\n\n\ndef get_sinusoidal_embeddings(\n    timesteps: jnp.ndarray,\n    embedding_dim: int,\n    freq_shift: float = 1,\n    min_timescale: float = 1,\n    max_timescale: float = 1.0e4,\n    flip_sin_to_cos: bool = False,\n    scale: float = 1.0,\n) -> jnp.ndarray:\n    \"\"\"Returns the positional encoding (same as Tensor2Tensor).\n    Args:\n        timesteps: a 1-D Tensor of N indices, one per batch element.\n        These may be fractional.\n        embedding_dim: The number of output channels.\n        min_timescale: The smallest time unit (should probably be 0.0).\n        max_timescale: The largest time unit.\n    Returns:\n        a Tensor of timing signals [N, num_channels]\n    \"\"\"\n    assert timesteps.ndim == 1, \"Timesteps should be a 1d-array\"\n    assert embedding_dim % 2 == 0, f\"Embedding dimension {embedding_dim} should be even\"\n    num_timescales = float(embedding_dim // 2)\n    log_timescale_increment = math.log(max_timescale / min_timescale) / (num_timescales - freq_shift)\n    inv_timescales = min_timescale * jnp.exp(jnp.arange(num_timescales, dtype=jnp.float32) * -log_timescale_increment)\n    emb = jnp.expand_dims(timesteps, 1) * jnp.expand_dims(inv_timescales, 0)\n\n    # scale embeddings\n    scaled_time = scale * emb\n\n    if flip_sin_to_cos:\n        signal = jnp.concatenate([jnp.cos(scaled_time), jnp.sin(scaled_time)], axis=1)\n    else:\n        signal = jnp.concatenate([jnp.sin(scaled_time), jnp.cos(scaled_time)], axis=1)\n    signal = jnp.reshape(signal, [jnp.shape(timesteps)[0], embedding_dim])\n    return signal\n\n\nclass FlaxTimestepEmbedding(nn.Module):\n    r\"\"\"\n    Time step Embedding Module. Learns embeddings for input time steps.\n\n    Args:\n        time_embed_dim (`int`, *optional*, defaults to `32`):\n                Time step embedding dimension\n        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):\n                Parameters `dtype`\n    \"\"\"\n    time_embed_dim: int = 32\n    dtype: jnp.dtype = jnp.float32\n\n    @nn.compact\n    def __call__(self, temb):\n        temb = nn.Dense(self.time_embed_dim, dtype=self.dtype, name=\"linear_1\")(temb)\n        temb = nn.silu(temb)\n        temb = nn.Dense(self.time_embed_dim, dtype=self.dtype, name=\"linear_2\")(temb)\n        return temb\n\n\nclass FlaxTimesteps(nn.Module):\n    r\"\"\"\n    Wrapper Module for sinusoidal Time step Embeddings as described in https://arxiv.org/abs/2006.11239\n\n    Args:\n        dim (`int`, *optional*, defaults to `32`):\n                Time step embedding dimension\n    \"\"\"\n    dim: int = 32\n    freq_shift: float = 1\n\n    @nn.compact\n    def __call__(self, timesteps):\n        return get_sinusoidal_embeddings(timesteps, embedding_dim=self.dim, freq_shift=self.freq_shift)\n", "path": "src/diffusers/models/embeddings_flax.py"}], "after_files": [{"content": "# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport math\n\nimport flax.linen as nn\nimport jax.numpy as jnp\n\n\ndef get_sinusoidal_embeddings(\n    timesteps: jnp.ndarray,\n    embedding_dim: int,\n    freq_shift: float = 1,\n    min_timescale: float = 1,\n    max_timescale: float = 1.0e4,\n    flip_sin_to_cos: bool = False,\n    scale: float = 1.0,\n) -> jnp.ndarray:\n    \"\"\"Returns the positional encoding (same as Tensor2Tensor).\n    Args:\n        timesteps: a 1-D Tensor of N indices, one per batch element.\n        These may be fractional.\n        embedding_dim: The number of output channels.\n        min_timescale: The smallest time unit (should probably be 0.0).\n        max_timescale: The largest time unit.\n    Returns:\n        a Tensor of timing signals [N, num_channels]\n    \"\"\"\n    assert timesteps.ndim == 1, \"Timesteps should be a 1d-array\"\n    assert embedding_dim % 2 == 0, f\"Embedding dimension {embedding_dim} should be even\"\n    num_timescales = float(embedding_dim // 2)\n    log_timescale_increment = math.log(max_timescale / min_timescale) / (num_timescales - freq_shift)\n    inv_timescales = min_timescale * jnp.exp(jnp.arange(num_timescales, dtype=jnp.float32) * -log_timescale_increment)\n    emb = jnp.expand_dims(timesteps, 1) * jnp.expand_dims(inv_timescales, 0)\n\n    # scale embeddings\n    scaled_time = scale * emb\n\n    if flip_sin_to_cos:\n        signal = jnp.concatenate([jnp.cos(scaled_time), jnp.sin(scaled_time)], axis=1)\n    else:\n        signal = jnp.concatenate([jnp.sin(scaled_time), jnp.cos(scaled_time)], axis=1)\n    signal = jnp.reshape(signal, [jnp.shape(timesteps)[0], embedding_dim])\n    return signal\n\n\nclass FlaxTimestepEmbedding(nn.Module):\n    r\"\"\"\n    Time step Embedding Module. Learns embeddings for input time steps.\n\n    Args:\n        time_embed_dim (`int`, *optional*, defaults to `32`):\n                Time step embedding dimension\n        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):\n                Parameters `dtype`\n    \"\"\"\n    time_embed_dim: int = 32\n    dtype: jnp.dtype = jnp.float32\n\n    @nn.compact\n    def __call__(self, temb):\n        temb = nn.Dense(self.time_embed_dim, dtype=self.dtype, name=\"linear_1\")(temb)\n        temb = nn.silu(temb)\n        temb = nn.Dense(self.time_embed_dim, dtype=self.dtype, name=\"linear_2\")(temb)\n        return temb\n\n\nclass FlaxTimesteps(nn.Module):\n    r\"\"\"\n    Wrapper Module for sinusoidal Time step Embeddings as described in https://arxiv.org/abs/2006.11239\n\n    Args:\n        dim (`int`, *optional*, defaults to `32`):\n                Time step embedding dimension\n    \"\"\"\n    dim: int = 32\n    freq_shift: float = 1\n\n    @nn.compact\n    def __call__(self, timesteps):\n        return get_sinusoidal_embeddings(\n            timesteps, embedding_dim=self.dim, freq_shift=self.freq_shift, flip_sin_to_cos=True\n        )\n", "path": "src/diffusers/models/embeddings_flax.py"}]}