We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
gumble_softmax
add mindspore gumble_softmax

- [gumble_softmax] #21042

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `ivy/functional/frontends/mindspore/ops/function/nn_func.py`
Content:
```
1 """Includes Mindspore Frontend functions listed in the TODO list
2 https://github.com/unifyai/ivy/issues/14951."""
3 
4 # local
5 import ivy
6 from ivy.func_wrapper import with_supported_dtypes
7 from ivy.functional.frontends.paddle.func_wrapper import to_ivy_arrays_and_back
8 
9 
10 def _broadcast_pooling_helper(x, pool_dims: str = "2d", name: str = "padding"):
11     dims = {"1d": 1, "2d": 2, "3d": 3}
12 
13     if isinstance(x, int):
14         return tuple([x for _ in range(dims[pool_dims])])
15 
16     if len(x) == 1:
17         return tuple([x[0] for _ in range(dims[pool_dims])])
18     elif len(x) == dims[pool_dims]:
19         return tuple(x)
20     elif len(x) != dims[pool_dims]:
21         raise ValueError(
22             f"`{name}` must either be a single int, "
23             f"or a tuple of {dims[pool_dims]} ints. "
24         )
25 
26 
27 @with_supported_dtypes(
28     {
29         "2.0.0 and below": (
30             "int8",
31             "int16",
32             "int32",
33             "int64",
34             "float16",
35             "float32",
36             "float64",
37         )
38     },
39     "mindspore",
40 )
41 @to_ivy_arrays_and_back
42 def dropout2d(input, p=0.5, training=True):
43     return ivy.dropout2d(input, p, training=training, data_format="NCHW")
44 
45 
46 @with_supported_dtypes({"2.0.0 and below": ("float16", "float32")}, "mindspore")
47 @to_ivy_arrays_and_back
48 def selu(input_x):
49     return ivy.selu(input_x)
50 
51 
52 @with_supported_dtypes({"2.0 and below": ("float16", "float32")}, "mindspore")
53 @to_ivy_arrays_and_back
54 def softsign(x):
55     return ivy.divide(x, ivy.add(1, ivy.abs(x)))
56 
57 
58 @with_supported_dtypes({"2.0.0 and below": ("float16", "float32")}, "mindspore")
59 @to_ivy_arrays_and_back
60 def log_softmax(input, axis=-1):
61     return ivy.log_softmax(input)
62 
63 
64 @with_supported_dtypes({"2.0 and below": ("float16", "float32")}, "mindspore")
65 @to_ivy_arrays_and_back
66 def kl_div(logits, labels, reduction="mean"):
67     """
68     Computes the Kullback-Leibler (KL) Divergence between the logits and the labels.
69 
70     Parameters:
71         logits (numpy array): The input logits array.
72         labels (numpy array): The label array which has the same shape as logits.
73         reduction (str): Specifies the reduction to be applied to the output.
74                          Its value must be one of 'none', 'mean', 'batchmean',
75                          or 'sum'. Default: 'mean'.
76 
77     Returns:
78         float or numpy array: If reduction is 'none', then output is
79         a numpy array and has the same shape as logits.
80                               Otherwise, it is a scalar (float).
81     """
82     assert ivy.shape(logits) == ivy.shape(
83         labels
84     ), "logits and labels must have the same shape."
85     L = labels * (ivy.log(labels) - logits)
86     if reduction == "none":
87         return L
88     elif reduction == "mean":
89         return ivy.mean(L)
90     elif reduction == "batchmean":
91         return ivy.mean(L, axis=0)
92     elif reduction == "sum":
93         return ivy.sum(L)
94     else:
95         raise ValueError(
96             "Invalid reduction mode. Supported values are 'none', 'mean', 'batchmean',"
97             " or 'sum'."
98         )
99 
100 
101 @with_supported_dtypes(
102     {
103         "2.0.0 and below": (
104             "int8",
105             "int16",
106             "int32",
107             "int64",
108             "float16",
109             "float32",
110             "float64",
111         )
112     },
113     "mindspore",
114 )
115 @to_ivy_arrays_and_back
116 def dropout3d(input, p=0.5, training=True):
117     return ivy.dropout3d(input, p, training=training, data_format="NCDHW")
118 
119 
120 @with_supported_dtypes(
121     {
122         "2.0.0 and below": (
123             "int8",
124             "int16",
125             "int32",
126             "int64",
127             "float16",
128             "float32",
129             "float64",
130         )
131     },
132     "mindspore",
133 )
134 @to_ivy_arrays_and_back
135 def interpolate(
136     input,
137     size=None,
138     scale_factor=None,
139     mode="nearest",
140     align_corners=False,
141     recompute_scale_factor=False,
142 ):
143     return ivy.interpolate(
144         input,
145         size=size,
146         scale_factor=scale_factor,
147         mode=mode,
148         align_corners=align_corners,
149         recompute_scale_factor=recompute_scale_factor,
150     )
151 
152 
153 @with_supported_dtypes(
154     {
155         "2.0 and below": (
156             "int8",
157             "int16",
158             "int32",
159             "int64",
160             "float16",
161             "float32",
162             "float64",
163         )
164     },
165     "mindspore",
166 )
167 @to_ivy_arrays_and_back
168 def pad(input, pad_width, mode="constant", constant_values=0):
169     return ivy.pad(input, pad_width, mode=mode, constant_values=constant_values)
170 
171 
172 @with_supported_dtypes(
173     {"2.0.0 and below": ("float16", "float32", "float64")}, "mindspore"
174 )
175 @to_ivy_arrays_and_back
176 def adaptive_avg_pool2d(input, output_size):
177     return ivy.adaptive_avg_pool2d(input, output_size)
178 
179 
180 @to_ivy_arrays_and_back
181 def avg_pool2d(
182     input,
183     kernel_size,
184     stride=None,
185     padding=0,
186     pad_mode=False,
187     count_include_pad=True,
188     divisor_override=None,
189 ):
190     # Figure out input dims N
191     input_rank = input.ndim
192 
193     if input_rank == 4:
194         # NCHW
195         data_format = "NCHW"
196 
197     kernel_size = _broadcast_pooling_helper(kernel_size, "2d", name="kernel_size")
198     stride = _broadcast_pooling_helper(stride, "2d", name="stride")
199     padding = _broadcast_pooling_helper(padding, "2d", name="padding")
200     kernel_pads = list(zip(kernel_size, padding))
201 
202     # Padding should be less than or equal to half of kernel size
203     if not all([pad <= kernel / 2 for kernel, pad in kernel_pads]):
204         raise ValueError(
205             "pad should be smaller than or equal to half of kernel size, "
206             f"but got padding={padding}, kernel_size={kernel_size}. "
207         )
208 
209     # Figure out padding string
210     if all([pad == ivy.ceil((kernel - 1) / 2) for kernel, pad in kernel_pads]):
211         padding_str = "SAME"
212     else:
213         padding_str = "VALID"
214 
215     return ivy.avg_pool2d(
216         input,
217         kernel_size,
218         stride,
219         padding_str,
220         data_format=data_format,
221         pad_mode=pad_mode,
222         count_include_pad=count_include_pad,
223         divisor_override=divisor_override,
224     )
225 
226 
227 @to_ivy_arrays_and_back
228 def flatten(input, order="C", *, start_dim=1, end_dim=-1):
229     return ivy.flatten(input, order=order, start_dim=start_dim, end_dim=end_dim)
230 
231 
232 @with_supported_dtypes(
233     {"2.0.0 and below": ("float16", "float32", "float64")},
234     "mindspore",
235 )
236 @to_ivy_arrays_and_back
237 def fast_gelu(input_x):
238     return (input_x / (1 + ivy.exp(-1.702 * ivy.abs(input_x)))) * ivy.exp(
239         0.851 * (input_x - ivy.abs(input_x))
240     )
241
```
--- 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()

```

diff --git a/ivy/functional/frontends/mindspore/ops/function/nn_func.py b/ivy/functional/frontends/mindspore/ops/function/nn_func.py
--- a/ivy/functional/frontends/mindspore/ops/function/nn_func.py
+++ b/ivy/functional/frontends/mindspore/ops/function/nn_func.py
@@ -238,3 +238,23 @@
     return (input_x / (1 + ivy.exp(-1.702 * ivy.abs(input_x)))) * ivy.exp(
         0.851 * (input_x - ivy.abs(input_x))
     )
+
+
+@with_supported_dtypes({"2.0.0 and below": ("float16", "float32")}, "mindspore")
+@to_ivy_arrays_and_back
+def gumbel_softmax(logits, tau=1, hard=False, dim=-1):
+    gumbels = -ivy.empty_like(logits).exponential().log()
+    gumbels = (logits + gumbels) / tau
+    y_soft = ivy.softmax(gumbels, axis=dim)
+
+    if hard:
+        indices = y_soft.max(axis=dim, keepdims=True)[1]
+        y_hard = ivy.zeros_like(logits)
+        updates = ivy.ones_like(indices)
+        y_hard = ivy.scatter_nd(indices, updates, reduction="replace", out=y_hard)
+
+        ret = y_hard - y_soft.stop_gradient(preserve_type=True) + y_soft
+    else:
+        ret = y_soft
+
+    return ret

{"golden_diff": "diff --git a/ivy/functional/frontends/mindspore/ops/function/nn_func.py b/ivy/functional/frontends/mindspore/ops/function/nn_func.py\n--- a/ivy/functional/frontends/mindspore/ops/function/nn_func.py\n+++ b/ivy/functional/frontends/mindspore/ops/function/nn_func.py\n@@ -238,3 +238,23 @@\n     return (input_x / (1 + ivy.exp(-1.702 * ivy.abs(input_x)))) * ivy.exp(\n         0.851 * (input_x - ivy.abs(input_x))\n     )\n+\n+\n+@with_supported_dtypes({\"2.0.0 and below\": (\"float16\", \"float32\")}, \"mindspore\")\n+@to_ivy_arrays_and_back\n+def gumbel_softmax(logits, tau=1, hard=False, dim=-1):\n+    gumbels = -ivy.empty_like(logits).exponential().log()\n+    gumbels = (logits + gumbels) / tau\n+    y_soft = ivy.softmax(gumbels, axis=dim)\n+\n+    if hard:\n+        indices = y_soft.max(axis=dim, keepdims=True)[1]\n+        y_hard = ivy.zeros_like(logits)\n+        updates = ivy.ones_like(indices)\n+        y_hard = ivy.scatter_nd(indices, updates, reduction=\"replace\", out=y_hard)\n+\n+        ret = y_hard - y_soft.stop_gradient(preserve_type=True) + y_soft\n+    else:\n+        ret = y_soft\n+\n+    return ret\n", "issue": "gumble_softmax\nadd mindspore gumble_softmax\r\n\r\n- [gumble_softmax] #21042\n", "before_files": [{"content": "\"\"\"Includes Mindspore Frontend functions listed in the TODO list\nhttps://github.com/unifyai/ivy/issues/14951.\"\"\"\n\n# local\nimport ivy\nfrom ivy.func_wrapper import with_supported_dtypes\nfrom ivy.functional.frontends.paddle.func_wrapper import to_ivy_arrays_and_back\n\n\ndef _broadcast_pooling_helper(x, pool_dims: str = \"2d\", name: str = \"padding\"):\n    dims = {\"1d\": 1, \"2d\": 2, \"3d\": 3}\n\n    if isinstance(x, int):\n        return tuple([x for _ in range(dims[pool_dims])])\n\n    if len(x) == 1:\n        return tuple([x[0] for _ in range(dims[pool_dims])])\n    elif len(x) == dims[pool_dims]:\n        return tuple(x)\n    elif len(x) != dims[pool_dims]:\n        raise ValueError(\n            f\"`{name}` must either be a single int, \"\n            f\"or a tuple of {dims[pool_dims]} ints. \"\n        )\n\n\n@with_supported_dtypes(\n    {\n        \"2.0.0 and below\": (\n            \"int8\",\n            \"int16\",\n            \"int32\",\n            \"int64\",\n            \"float16\",\n            \"float32\",\n            \"float64\",\n        )\n    },\n    \"mindspore\",\n)\n@to_ivy_arrays_and_back\ndef dropout2d(input, p=0.5, training=True):\n    return ivy.dropout2d(input, p, training=training, data_format=\"NCHW\")\n\n\n@with_supported_dtypes({\"2.0.0 and below\": (\"float16\", \"float32\")}, \"mindspore\")\n@to_ivy_arrays_and_back\ndef selu(input_x):\n    return ivy.selu(input_x)\n\n\n@with_supported_dtypes({\"2.0 and below\": (\"float16\", \"float32\")}, \"mindspore\")\n@to_ivy_arrays_and_back\ndef softsign(x):\n    return ivy.divide(x, ivy.add(1, ivy.abs(x)))\n\n\n@with_supported_dtypes({\"2.0.0 and below\": (\"float16\", \"float32\")}, \"mindspore\")\n@to_ivy_arrays_and_back\ndef log_softmax(input, axis=-1):\n    return ivy.log_softmax(input)\n\n\n@with_supported_dtypes({\"2.0 and below\": (\"float16\", \"float32\")}, \"mindspore\")\n@to_ivy_arrays_and_back\ndef kl_div(logits, labels, reduction=\"mean\"):\n    \"\"\"\n    Computes the Kullback-Leibler (KL) Divergence between the logits and the labels.\n\n    Parameters:\n        logits (numpy array): The input logits array.\n        labels (numpy array): The label array which has the same shape as logits.\n        reduction (str): Specifies the reduction to be applied to the output.\n                         Its value must be one of 'none', 'mean', 'batchmean',\n                         or 'sum'. Default: 'mean'.\n\n    Returns:\n        float or numpy array: If reduction is 'none', then output is\n        a numpy array and has the same shape as logits.\n                              Otherwise, it is a scalar (float).\n    \"\"\"\n    assert ivy.shape(logits) == ivy.shape(\n        labels\n    ), \"logits and labels must have the same shape.\"\n    L = labels * (ivy.log(labels) - logits)\n    if reduction == \"none\":\n        return L\n    elif reduction == \"mean\":\n        return ivy.mean(L)\n    elif reduction == \"batchmean\":\n        return ivy.mean(L, axis=0)\n    elif reduction == \"sum\":\n        return ivy.sum(L)\n    else:\n        raise ValueError(\n            \"Invalid reduction mode. Supported values are 'none', 'mean', 'batchmean',\"\n            \" or 'sum'.\"\n        )\n\n\n@with_supported_dtypes(\n    {\n        \"2.0.0 and below\": (\n            \"int8\",\n            \"int16\",\n            \"int32\",\n            \"int64\",\n            \"float16\",\n            \"float32\",\n            \"float64\",\n        )\n    },\n    \"mindspore\",\n)\n@to_ivy_arrays_and_back\ndef dropout3d(input, p=0.5, training=True):\n    return ivy.dropout3d(input, p, training=training, data_format=\"NCDHW\")\n\n\n@with_supported_dtypes(\n    {\n        \"2.0.0 and below\": (\n            \"int8\",\n            \"int16\",\n            \"int32\",\n            \"int64\",\n            \"float16\",\n            \"float32\",\n            \"float64\",\n        )\n    },\n    \"mindspore\",\n)\n@to_ivy_arrays_and_back\ndef interpolate(\n    input,\n    size=None,\n    scale_factor=None,\n    mode=\"nearest\",\n    align_corners=False,\n    recompute_scale_factor=False,\n):\n    return ivy.interpolate(\n        input,\n        size=size,\n        scale_factor=scale_factor,\n        mode=mode,\n        align_corners=align_corners,\n        recompute_scale_factor=recompute_scale_factor,\n    )\n\n\n@with_supported_dtypes(\n    {\n        \"2.0 and below\": (\n            \"int8\",\n            \"int16\",\n            \"int32\",\n            \"int64\",\n            \"float16\",\n            \"float32\",\n            \"float64\",\n        )\n    },\n    \"mindspore\",\n)\n@to_ivy_arrays_and_back\ndef pad(input, pad_width, mode=\"constant\", constant_values=0):\n    return ivy.pad(input, pad_width, mode=mode, constant_values=constant_values)\n\n\n@with_supported_dtypes(\n    {\"2.0.0 and below\": (\"float16\", \"float32\", \"float64\")}, \"mindspore\"\n)\n@to_ivy_arrays_and_back\ndef adaptive_avg_pool2d(input, output_size):\n    return ivy.adaptive_avg_pool2d(input, output_size)\n\n\n@to_ivy_arrays_and_back\ndef avg_pool2d(\n    input,\n    kernel_size,\n    stride=None,\n    padding=0,\n    pad_mode=False,\n    count_include_pad=True,\n    divisor_override=None,\n):\n    # Figure out input dims N\n    input_rank = input.ndim\n\n    if input_rank == 4:\n        # NCHW\n        data_format = \"NCHW\"\n\n    kernel_size = _broadcast_pooling_helper(kernel_size, \"2d\", name=\"kernel_size\")\n    stride = _broadcast_pooling_helper(stride, \"2d\", name=\"stride\")\n    padding = _broadcast_pooling_helper(padding, \"2d\", name=\"padding\")\n    kernel_pads = list(zip(kernel_size, padding))\n\n    # Padding should be less than or equal to half of kernel size\n    if not all([pad <= kernel / 2 for kernel, pad in kernel_pads]):\n        raise ValueError(\n            \"pad should be smaller than or equal to half of kernel size, \"\n            f\"but got padding={padding}, kernel_size={kernel_size}. \"\n        )\n\n    # Figure out padding string\n    if all([pad == ivy.ceil((kernel - 1) / 2) for kernel, pad in kernel_pads]):\n        padding_str = \"SAME\"\n    else:\n        padding_str = \"VALID\"\n\n    return ivy.avg_pool2d(\n        input,\n        kernel_size,\n        stride,\n        padding_str,\n        data_format=data_format,\n        pad_mode=pad_mode,\n        count_include_pad=count_include_pad,\n        divisor_override=divisor_override,\n    )\n\n\n@to_ivy_arrays_and_back\ndef flatten(input, order=\"C\", *, start_dim=1, end_dim=-1):\n    return ivy.flatten(input, order=order, start_dim=start_dim, end_dim=end_dim)\n\n\n@with_supported_dtypes(\n    {\"2.0.0 and below\": (\"float16\", \"float32\", \"float64\")},\n    \"mindspore\",\n)\n@to_ivy_arrays_and_back\ndef fast_gelu(input_x):\n    return (input_x / (1 + ivy.exp(-1.702 * ivy.abs(input_x)))) * ivy.exp(\n        0.851 * (input_x - ivy.abs(input_x))\n    )\n", "path": "ivy/functional/frontends/mindspore/ops/function/nn_func.py"}], "after_files": [{"content": "\"\"\"Includes Mindspore Frontend functions listed in the TODO list\nhttps://github.com/unifyai/ivy/issues/14951.\"\"\"\n\n# local\nimport ivy\nfrom ivy.func_wrapper import with_supported_dtypes\nfrom ivy.functional.frontends.paddle.func_wrapper import to_ivy_arrays_and_back\n\n\ndef _broadcast_pooling_helper(x, pool_dims: str = \"2d\", name: str = \"padding\"):\n    dims = {\"1d\": 1, \"2d\": 2, \"3d\": 3}\n\n    if isinstance(x, int):\n        return tuple([x for _ in range(dims[pool_dims])])\n\n    if len(x) == 1:\n        return tuple([x[0] for _ in range(dims[pool_dims])])\n    elif len(x) == dims[pool_dims]:\n        return tuple(x)\n    elif len(x) != dims[pool_dims]:\n        raise ValueError(\n            f\"`{name}` must either be a single int, \"\n            f\"or a tuple of {dims[pool_dims]} ints. \"\n        )\n\n\n@with_supported_dtypes(\n    {\n        \"2.0.0 and below\": (\n            \"int8\",\n            \"int16\",\n            \"int32\",\n            \"int64\",\n            \"float16\",\n            \"float32\",\n            \"float64\",\n        )\n    },\n    \"mindspore\",\n)\n@to_ivy_arrays_and_back\ndef dropout2d(input, p=0.5, training=True):\n    return ivy.dropout2d(input, p, training=training, data_format=\"NCHW\")\n\n\n@with_supported_dtypes({\"2.0.0 and below\": (\"float16\", \"float32\")}, \"mindspore\")\n@to_ivy_arrays_and_back\ndef selu(input_x):\n    return ivy.selu(input_x)\n\n\n@with_supported_dtypes({\"2.0 and below\": (\"float16\", \"float32\")}, \"mindspore\")\n@to_ivy_arrays_and_back\ndef softsign(x):\n    return ivy.divide(x, ivy.add(1, ivy.abs(x)))\n\n\n@with_supported_dtypes({\"2.0.0 and below\": (\"float16\", \"float32\")}, \"mindspore\")\n@to_ivy_arrays_and_back\ndef log_softmax(input, axis=-1):\n    return ivy.log_softmax(input)\n\n\n@with_supported_dtypes({\"2.0 and below\": (\"float16\", \"float32\")}, \"mindspore\")\n@to_ivy_arrays_and_back\ndef kl_div(logits, labels, reduction=\"mean\"):\n    \"\"\"\n    Computes the Kullback-Leibler (KL) Divergence between the logits and the labels.\n\n    Parameters:\n        logits (numpy array): The input logits array.\n        labels (numpy array): The label array which has the same shape as logits.\n        reduction (str): Specifies the reduction to be applied to the output.\n                         Its value must be one of 'none', 'mean', 'batchmean',\n                         or 'sum'. Default: 'mean'.\n\n    Returns:\n        float or numpy array: If reduction is 'none', then output is\n        a numpy array and has the same shape as logits.\n                              Otherwise, it is a scalar (float).\n    \"\"\"\n    assert ivy.shape(logits) == ivy.shape(\n        labels\n    ), \"logits and labels must have the same shape.\"\n    L = labels * (ivy.log(labels) - logits)\n    if reduction == \"none\":\n        return L\n    elif reduction == \"mean\":\n        return ivy.mean(L)\n    elif reduction == \"batchmean\":\n        return ivy.mean(L, axis=0)\n    elif reduction == \"sum\":\n        return ivy.sum(L)\n    else:\n        raise ValueError(\n            \"Invalid reduction mode. Supported values are 'none', 'mean', 'batchmean',\"\n            \" or 'sum'.\"\n        )\n\n\n@with_supported_dtypes(\n    {\n        \"2.0.0 and below\": (\n            \"int8\",\n            \"int16\",\n            \"int32\",\n            \"int64\",\n            \"float16\",\n            \"float32\",\n            \"float64\",\n        )\n    },\n    \"mindspore\",\n)\n@to_ivy_arrays_and_back\ndef dropout3d(input, p=0.5, training=True):\n    return ivy.dropout3d(input, p, training=training, data_format=\"NCDHW\")\n\n\n@with_supported_dtypes(\n    {\n        \"2.0.0 and below\": (\n            \"int8\",\n            \"int16\",\n            \"int32\",\n            \"int64\",\n            \"float16\",\n            \"float32\",\n            \"float64\",\n        )\n    },\n    \"mindspore\",\n)\n@to_ivy_arrays_and_back\ndef interpolate(\n    input,\n    size=None,\n    scale_factor=None,\n    mode=\"nearest\",\n    align_corners=False,\n    recompute_scale_factor=False,\n):\n    return ivy.interpolate(\n        input,\n        size=size,\n        scale_factor=scale_factor,\n        mode=mode,\n        align_corners=align_corners,\n        recompute_scale_factor=recompute_scale_factor,\n    )\n\n\n@with_supported_dtypes(\n    {\n        \"2.0 and below\": (\n            \"int8\",\n            \"int16\",\n            \"int32\",\n            \"int64\",\n            \"float16\",\n            \"float32\",\n            \"float64\",\n        )\n    },\n    \"mindspore\",\n)\n@to_ivy_arrays_and_back\ndef pad(input, pad_width, mode=\"constant\", constant_values=0):\n    return ivy.pad(input, pad_width, mode=mode, constant_values=constant_values)\n\n\n@with_supported_dtypes(\n    {\"2.0.0 and below\": (\"float16\", \"float32\", \"float64\")}, \"mindspore\"\n)\n@to_ivy_arrays_and_back\ndef adaptive_avg_pool2d(input, output_size):\n    return ivy.adaptive_avg_pool2d(input, output_size)\n\n\n@to_ivy_arrays_and_back\ndef avg_pool2d(\n    input,\n    kernel_size,\n    stride=None,\n    padding=0,\n    pad_mode=False,\n    count_include_pad=True,\n    divisor_override=None,\n):\n    # Figure out input dims N\n    input_rank = input.ndim\n\n    if input_rank == 4:\n        # NCHW\n        data_format = \"NCHW\"\n\n    kernel_size = _broadcast_pooling_helper(kernel_size, \"2d\", name=\"kernel_size\")\n    stride = _broadcast_pooling_helper(stride, \"2d\", name=\"stride\")\n    padding = _broadcast_pooling_helper(padding, \"2d\", name=\"padding\")\n    kernel_pads = list(zip(kernel_size, padding))\n\n    # Padding should be less than or equal to half of kernel size\n    if not all([pad <= kernel / 2 for kernel, pad in kernel_pads]):\n        raise ValueError(\n            \"pad should be smaller than or equal to half of kernel size, \"\n            f\"but got padding={padding}, kernel_size={kernel_size}. \"\n        )\n\n    # Figure out padding string\n    if all([pad == ivy.ceil((kernel - 1) / 2) for kernel, pad in kernel_pads]):\n        padding_str = \"SAME\"\n    else:\n        padding_str = \"VALID\"\n\n    return ivy.avg_pool2d(\n        input,\n        kernel_size,\n        stride,\n        padding_str,\n        data_format=data_format,\n        pad_mode=pad_mode,\n        count_include_pad=count_include_pad,\n        divisor_override=divisor_override,\n    )\n\n\n@to_ivy_arrays_and_back\ndef flatten(input, order=\"C\", *, start_dim=1, end_dim=-1):\n    return ivy.flatten(input, order=order, start_dim=start_dim, end_dim=end_dim)\n\n\n@with_supported_dtypes(\n    {\"2.0.0 and below\": (\"float16\", \"float32\", \"float64\")},\n    \"mindspore\",\n)\n@to_ivy_arrays_and_back\ndef fast_gelu(input_x):\n    return (input_x / (1 + ivy.exp(-1.702 * ivy.abs(input_x)))) * ivy.exp(\n        0.851 * (input_x - ivy.abs(input_x))\n    )\n\n\n@with_supported_dtypes({\"2.0.0 and below\": (\"float16\", \"float32\")}, \"mindspore\")\n@to_ivy_arrays_and_back\ndef gumbel_softmax(logits, tau=1, hard=False, dim=-1):\n    gumbels = -ivy.empty_like(logits).exponential().log()\n    gumbels = (logits + gumbels) / tau\n    y_soft = ivy.softmax(gumbels, axis=dim)\n\n    if hard:\n        indices = y_soft.max(axis=dim, keepdims=True)[1]\n        y_hard = ivy.zeros_like(logits)\n        updates = ivy.ones_like(indices)\n        y_hard = ivy.scatter_nd(indices, updates, reduction=\"replace\", out=y_hard)\n\n        ret = y_hard - y_soft.stop_gradient(preserve_type=True) + y_soft\n    else:\n        ret = y_soft\n\n    return ret\n", "path": "ivy/functional/frontends/mindspore/ops/function/nn_func.py"}]}