We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Implement log(cosh()) loss function
log-cosh is a doubly differentiable alternative to the huber loss. A naive implementation is prone to overflow (since cosh has an e^x term), so I think it'd be a useful addition to the library. Plus, it's implemented in other libraries, such as [TensorFlow](https://www.tensorflow.org/api_docs/python/tf/keras/losses/log_cosh).

If this sounds like a relevant addition, I'd be happy to contribute it!

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `optax/_src/loss.py`
Content:
```
1 # Copyright 2019 DeepMind Technologies Limited. 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 # ==============================================================================
15 """Standard losses used in optimisation.
16 
17 We provide implementations of the most canonical losses used in deep
18 learning. These operate transparently on batches, and do not perform any
19 reduction over the batch dimensions, leaving it to the user to, for instance,
20 mean or sum losses across batch dimensions.
21 """
22 
23 from typing import Optional
24 
25 import chex
26 import jax
27 import jax.numpy as jnp
28 
29 from optax._src import utils
30 
31 
32 def l2_loss(
33     predictions: chex.Array,
34     targets: Optional[chex.Array] = None,
35 ) -> chex.Array:
36   """Calculates the L2 loss for a set of predictions.
37 
38   Note: the 0.5 term is standard in "Pattern Recognition and Machine Learning"
39   by Bishop, but not "The Elements of Statistical Learning" by Tibshirani.
40 
41   References:
42     [Chris Bishop, 2006](https://bit.ly/3eeP0ga)
43 
44   Args:
45     predictions: a vector of arbitrary shape.
46     targets: a vector of shape compatible with predictions; if not provides
47       then it is assumed to be zero.
48 
49   Returns:
50     the squared error loss.
51   """
52   chex.assert_type([predictions], float)
53   errors = (predictions - targets) if (targets is not None) else predictions
54   return 0.5 * (errors)**2
55 
56 
57 def huber_loss(
58     predictions: chex.Array,
59     targets: Optional[chex.Array] = None,
60     delta: float = 1.) -> chex.Array:
61   """Huber loss, similar to L2 loss close to zero, L1 loss away from zero.
62 
63   If gradient descent is applied to the `huber loss`, it is equivalent to
64   clipping gradients of an `l2_loss` to `[-delta, delta]` in the backward pass.
65 
66   References:
67     [Huber, 1964](www.projecteuclid.org/download/pdf_1/euclid.aoms/1177703732)
68 
69   Args:
70     predictions: a vector of arbitrary shape.
71     targets: a vector of shape compatible with predictions; if not provides
72       then it is assumed to be zero.
73     delta: the bounds for the huber loss transformation, defaults at 1.
74 
75   Returns:
76     a vector of same shape of `x`.
77   """
78   chex.assert_type([predictions], float)
79   errors = (predictions - targets) if (targets is not None) else predictions
80   # 0.5 * err^2                  if |err| <= d
81   # 0.5 * d^2 + d * (|err| - d)  if |err| > d
82   abs_errors = jnp.abs(errors)
83   quadratic = jnp.minimum(abs_errors, delta)
84   # Same as max(abs_x - delta, 0) but avoids potentially doubling gradient.
85   linear = abs_errors - quadratic
86   return 0.5 * quadratic ** 2 + delta * linear
87 
88 
89 def smooth_labels(
90     labels: chex.Array,
91     alpha: float,
92 ) -> jnp.ndarray:
93   """Apply label smoothing.
94 
95   Label smoothing is often used in combination with a cross-entropy loss.
96   Smoothed labels favour small logit gaps, and it has been shown that this can
97   provide better model calibration by preventing overconfident predictions.
98 
99   References:
100     [Müller et al, 2019](https://arxiv.org/pdf/1906.02629.pdf)
101 
102   Args:
103     labels: one hot labels to be smoothed.
104     alpha: the smoothing factor, the greedy category with be assigned
105       probability `(1-alpha) + alpha / num_categories`
106 
107   Returns:
108     a smoothed version of the one hot input labels.
109 
110   """
111   chex.assert_type([labels], float)
112   num_categories = labels.shape[-1]
113   return (1.0 - alpha) * labels + alpha / num_categories
114 
115 
116 def sigmoid_binary_cross_entropy(logits, labels):
117   """Computes sigmoid cross entropy given logits and multiple class labels.
118 
119   Measures the probability error in discrete classification tasks in which
120   each class is an independent binary prediction and different classes are
121   not mutually exclusive. This may be used for multilabel image classification
122   for instance a model may predict that an image contains both a cat and a dog.
123 
124   References:
125     [Goodfellow et al, 2016](http://www.deeplearningbook.org/contents/prob.html)
126 
127   Args:
128     logits: unnormalized log probabilities.
129     labels: the probability for that class.
130 
131   Returns:
132     a sigmoid cross entropy loss.
133   """
134   chex.assert_equal_shape([logits, labels])
135   chex.assert_type([logits, labels], float)
136   log_p = jax.nn.log_sigmoid(logits)
137   # log(1 - sigmoid(x)) = log_sigmoid(-x), the latter more numerically stable
138   log_not_p = jax.nn.log_sigmoid(-logits)
139   return -labels * log_p - (1. - labels) * log_not_p
140 
141 
142 def softmax_cross_entropy(
143     logits: chex.Array,
144     labels: chex.Array,
145 ) -> chex.Array:
146   """Computes the softmax cross entropy between sets of logits and labels.
147 
148   Measures the probability error in discrete classification tasks in which
149   the classes are mutually exclusive (each entry is in exactly one class).
150   For example, each CIFAR-10 image is labeled with one and only one label:
151   an image can be a dog or a truck, but not both.
152 
153   References:
154     [Goodfellow et al, 2016](http://www.deeplearningbook.org/contents/prob.html)
155 
156   Args:
157     logits: unnormalized log probabilities.
158     labels: a valid probability distribution (non-negative, sum to 1), e.g a
159       one hot encoding of which class is the correct one for each input.
160 
161   Returns:
162     the cross entropy loss.
163   """
164   chex.assert_equal_shape([logits, labels])
165   chex.assert_type([logits, labels], float)
166   return -jnp.sum(labels * jax.nn.log_softmax(logits, axis=-1), axis=-1)
167 
168 
169 def cosine_similarity(
170     predictions: chex.Array,
171     targets: chex.Array,
172     epsilon: float = 0.,
173 ) -> chex.Array:
174   r"""Computes the cosine similarity between targets and predictions.
175 
176   The cosine **similarity** is a measure of similarity between vectors defined
177   as the cosine of the angle between them, which is also the inner product of
178   those vectors normalized to have unit norm.
179 
180   References:
181     [Wikipedia, 2021](https://en.wikipedia.org/wiki/Cosine_similarity)
182 
183   Args:
184     predictions: The predicted vector.
185     targets: Ground truth target vector.
186     epsilon: minimum norm for terms in the denominator of the cosine similarity.
187 
188   Returns:
189     cosine similarity values.
190   """
191   chex.assert_equal_shape([targets, predictions])
192   chex.assert_type([targets, predictions], float)
193   # vectorize norm fn, to treat all dimensions except the last as batch dims.
194   batched_norm_fn = jnp.vectorize(
195       utils.safe_norm, signature='(k)->()', excluded={1})
196   # normalise the last dimension of targets and predictions.
197   unit_targets = targets / jnp.expand_dims(
198       batched_norm_fn(targets, epsilon), axis=-1)
199   unit_predictions = predictions / jnp.expand_dims(
200       batched_norm_fn(predictions, epsilon), axis=-1)
201   # return cosine similarity.
202   return jnp.sum(unit_targets * unit_predictions, axis=-1)
203 
204 
205 def cosine_distance(
206     predictions: chex.Array,
207     targets: chex.Array,
208     epsilon: float = 0.,
209 ) -> chex.Array:
210   r"""Computes the cosine distance between targets and predictions.
211 
212   The cosine **distance**, implemented here, measures the **dissimilarity**
213   of two vectors as the opposite of cosine **similarity**: `1 - cos(\theta)`.
214 
215   References:
216     [Wikipedia, 2021](https://en.wikipedia.org/wiki/Cosine_similarity)
217 
218   Args:
219     predictions: The predicted vector.
220     targets: Ground truth target vector.
221     epsilon: minimum norm for terms in the denominator of the cosine similarity.
222 
223   Returns:
224     cosine similarity values.
225   """
226   chex.assert_equal_shape([targets, predictions])
227   chex.assert_type([targets, predictions], float)
228   # cosine distance = 1 - cosine similarity.
229   return 1. - cosine_similarity(predictions, targets, epsilon)
230
```
--- 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/optax/_src/loss.py b/optax/_src/loss.py
--- a/optax/_src/loss.py
+++ b/optax/_src/loss.py
@@ -227,3 +227,29 @@
   chex.assert_type([targets, predictions], float)
   # cosine distance = 1 - cosine similarity.
   return 1. - cosine_similarity(predictions, targets, epsilon)
+
+
+def log_cosh(
+    predictions: chex.Array,
+    targets: Optional[chex.Array] = None,
+) -> chex.Array:
+  """Calculates the log-cosh loss for a set of predictions.
+
+  log(cosh(x)) is approximately `(x**2) / 2` for small x and `abs(x) - log(2)`
+  for large x.  It is a twice differentiable alternative to the Huber loss.
+
+  References:
+    [Chen et al, 2019](https://openreview.net/pdf?id=rkglvsC9Ym)
+
+  Args:
+    predictions: a vector of arbitrary shape.
+    targets: a vector of shape compatible with predictions; if not provided
+      then it is assumed to be zero.
+
+  Returns:
+    the log-cosh loss.
+  """
+  chex.assert_type([predictions], float)
+  errors = (predictions - targets) if (targets is not None) else predictions
+  # log(cosh(x)) = log((exp(x) + exp(-x))/2) = log(exp(x) + exp(-x)) - log(2)
+  return jnp.logaddexp(errors, -errors) - jnp.log(2.0).astype(errors.dtype)

{"golden_diff": "diff --git a/optax/_src/loss.py b/optax/_src/loss.py\n--- a/optax/_src/loss.py\n+++ b/optax/_src/loss.py\n@@ -227,3 +227,29 @@\n   chex.assert_type([targets, predictions], float)\n   # cosine distance = 1 - cosine similarity.\n   return 1. - cosine_similarity(predictions, targets, epsilon)\n+\n+\n+def log_cosh(\n+    predictions: chex.Array,\n+    targets: Optional[chex.Array] = None,\n+) -> chex.Array:\n+  \"\"\"Calculates the log-cosh loss for a set of predictions.\n+\n+  log(cosh(x)) is approximately `(x**2) / 2` for small x and `abs(x) - log(2)`\n+  for large x.  It is a twice differentiable alternative to the Huber loss.\n+\n+  References:\n+    [Chen et al, 2019](https://openreview.net/pdf?id=rkglvsC9Ym)\n+\n+  Args:\n+    predictions: a vector of arbitrary shape.\n+    targets: a vector of shape compatible with predictions; if not provided\n+      then it is assumed to be zero.\n+\n+  Returns:\n+    the log-cosh loss.\n+  \"\"\"\n+  chex.assert_type([predictions], float)\n+  errors = (predictions - targets) if (targets is not None) else predictions\n+  # log(cosh(x)) = log((exp(x) + exp(-x))/2) = log(exp(x) + exp(-x)) - log(2)\n+  return jnp.logaddexp(errors, -errors) - jnp.log(2.0).astype(errors.dtype)\n", "issue": "Implement log(cosh()) loss function\nlog-cosh is a doubly differentiable alternative to the huber loss. A naive implementation is prone to overflow (since cosh has an e^x term), so I think it'd be a useful addition to the library. Plus, it's implemented in other libraries, such as [TensorFlow](https://www.tensorflow.org/api_docs/python/tf/keras/losses/log_cosh).\r\n\r\nIf this sounds like a relevant addition, I'd be happy to contribute it!\n", "before_files": [{"content": "# Copyright 2019 DeepMind Technologies Limited. 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.\n# ==============================================================================\n\"\"\"Standard losses used in optimisation.\n\nWe provide implementations of the most canonical losses used in deep\nlearning. These operate transparently on batches, and do not perform any\nreduction over the batch dimensions, leaving it to the user to, for instance,\nmean or sum losses across batch dimensions.\n\"\"\"\n\nfrom typing import Optional\n\nimport chex\nimport jax\nimport jax.numpy as jnp\n\nfrom optax._src import utils\n\n\ndef l2_loss(\n    predictions: chex.Array,\n    targets: Optional[chex.Array] = None,\n) -> chex.Array:\n  \"\"\"Calculates the L2 loss for a set of predictions.\n\n  Note: the 0.5 term is standard in \"Pattern Recognition and Machine Learning\"\n  by Bishop, but not \"The Elements of Statistical Learning\" by Tibshirani.\n\n  References:\n    [Chris Bishop, 2006](https://bit.ly/3eeP0ga)\n\n  Args:\n    predictions: a vector of arbitrary shape.\n    targets: a vector of shape compatible with predictions; if not provides\n      then it is assumed to be zero.\n\n  Returns:\n    the squared error loss.\n  \"\"\"\n  chex.assert_type([predictions], float)\n  errors = (predictions - targets) if (targets is not None) else predictions\n  return 0.5 * (errors)**2\n\n\ndef huber_loss(\n    predictions: chex.Array,\n    targets: Optional[chex.Array] = None,\n    delta: float = 1.) -> chex.Array:\n  \"\"\"Huber loss, similar to L2 loss close to zero, L1 loss away from zero.\n\n  If gradient descent is applied to the `huber loss`, it is equivalent to\n  clipping gradients of an `l2_loss` to `[-delta, delta]` in the backward pass.\n\n  References:\n    [Huber, 1964](www.projecteuclid.org/download/pdf_1/euclid.aoms/1177703732)\n\n  Args:\n    predictions: a vector of arbitrary shape.\n    targets: a vector of shape compatible with predictions; if not provides\n      then it is assumed to be zero.\n    delta: the bounds for the huber loss transformation, defaults at 1.\n\n  Returns:\n    a vector of same shape of `x`.\n  \"\"\"\n  chex.assert_type([predictions], float)\n  errors = (predictions - targets) if (targets is not None) else predictions\n  # 0.5 * err^2                  if |err| <= d\n  # 0.5 * d^2 + d * (|err| - d)  if |err| > d\n  abs_errors = jnp.abs(errors)\n  quadratic = jnp.minimum(abs_errors, delta)\n  # Same as max(abs_x - delta, 0) but avoids potentially doubling gradient.\n  linear = abs_errors - quadratic\n  return 0.5 * quadratic ** 2 + delta * linear\n\n\ndef smooth_labels(\n    labels: chex.Array,\n    alpha: float,\n) -> jnp.ndarray:\n  \"\"\"Apply label smoothing.\n\n  Label smoothing is often used in combination with a cross-entropy loss.\n  Smoothed labels favour small logit gaps, and it has been shown that this can\n  provide better model calibration by preventing overconfident predictions.\n\n  References:\n    [M\u00fcller et al, 2019](https://arxiv.org/pdf/1906.02629.pdf)\n\n  Args:\n    labels: one hot labels to be smoothed.\n    alpha: the smoothing factor, the greedy category with be assigned\n      probability `(1-alpha) + alpha / num_categories`\n\n  Returns:\n    a smoothed version of the one hot input labels.\n\n  \"\"\"\n  chex.assert_type([labels], float)\n  num_categories = labels.shape[-1]\n  return (1.0 - alpha) * labels + alpha / num_categories\n\n\ndef sigmoid_binary_cross_entropy(logits, labels):\n  \"\"\"Computes sigmoid cross entropy given logits and multiple class labels.\n\n  Measures the probability error in discrete classification tasks in which\n  each class is an independent binary prediction and different classes are\n  not mutually exclusive. This may be used for multilabel image classification\n  for instance a model may predict that an image contains both a cat and a dog.\n\n  References:\n    [Goodfellow et al, 2016](http://www.deeplearningbook.org/contents/prob.html)\n\n  Args:\n    logits: unnormalized log probabilities.\n    labels: the probability for that class.\n\n  Returns:\n    a sigmoid cross entropy loss.\n  \"\"\"\n  chex.assert_equal_shape([logits, labels])\n  chex.assert_type([logits, labels], float)\n  log_p = jax.nn.log_sigmoid(logits)\n  # log(1 - sigmoid(x)) = log_sigmoid(-x), the latter more numerically stable\n  log_not_p = jax.nn.log_sigmoid(-logits)\n  return -labels * log_p - (1. - labels) * log_not_p\n\n\ndef softmax_cross_entropy(\n    logits: chex.Array,\n    labels: chex.Array,\n) -> chex.Array:\n  \"\"\"Computes the softmax cross entropy between sets of logits and labels.\n\n  Measures the probability error in discrete classification tasks in which\n  the classes are mutually exclusive (each entry is in exactly one class).\n  For example, each CIFAR-10 image is labeled with one and only one label:\n  an image can be a dog or a truck, but not both.\n\n  References:\n    [Goodfellow et al, 2016](http://www.deeplearningbook.org/contents/prob.html)\n\n  Args:\n    logits: unnormalized log probabilities.\n    labels: a valid probability distribution (non-negative, sum to 1), e.g a\n      one hot encoding of which class is the correct one for each input.\n\n  Returns:\n    the cross entropy loss.\n  \"\"\"\n  chex.assert_equal_shape([logits, labels])\n  chex.assert_type([logits, labels], float)\n  return -jnp.sum(labels * jax.nn.log_softmax(logits, axis=-1), axis=-1)\n\n\ndef cosine_similarity(\n    predictions: chex.Array,\n    targets: chex.Array,\n    epsilon: float = 0.,\n) -> chex.Array:\n  r\"\"\"Computes the cosine similarity between targets and predictions.\n\n  The cosine **similarity** is a measure of similarity between vectors defined\n  as the cosine of the angle between them, which is also the inner product of\n  those vectors normalized to have unit norm.\n\n  References:\n    [Wikipedia, 2021](https://en.wikipedia.org/wiki/Cosine_similarity)\n\n  Args:\n    predictions: The predicted vector.\n    targets: Ground truth target vector.\n    epsilon: minimum norm for terms in the denominator of the cosine similarity.\n\n  Returns:\n    cosine similarity values.\n  \"\"\"\n  chex.assert_equal_shape([targets, predictions])\n  chex.assert_type([targets, predictions], float)\n  # vectorize norm fn, to treat all dimensions except the last as batch dims.\n  batched_norm_fn = jnp.vectorize(\n      utils.safe_norm, signature='(k)->()', excluded={1})\n  # normalise the last dimension of targets and predictions.\n  unit_targets = targets / jnp.expand_dims(\n      batched_norm_fn(targets, epsilon), axis=-1)\n  unit_predictions = predictions / jnp.expand_dims(\n      batched_norm_fn(predictions, epsilon), axis=-1)\n  # return cosine similarity.\n  return jnp.sum(unit_targets * unit_predictions, axis=-1)\n\n\ndef cosine_distance(\n    predictions: chex.Array,\n    targets: chex.Array,\n    epsilon: float = 0.,\n) -> chex.Array:\n  r\"\"\"Computes the cosine distance between targets and predictions.\n\n  The cosine **distance**, implemented here, measures the **dissimilarity**\n  of two vectors as the opposite of cosine **similarity**: `1 - cos(\\theta)`.\n\n  References:\n    [Wikipedia, 2021](https://en.wikipedia.org/wiki/Cosine_similarity)\n\n  Args:\n    predictions: The predicted vector.\n    targets: Ground truth target vector.\n    epsilon: minimum norm for terms in the denominator of the cosine similarity.\n\n  Returns:\n    cosine similarity values.\n  \"\"\"\n  chex.assert_equal_shape([targets, predictions])\n  chex.assert_type([targets, predictions], float)\n  # cosine distance = 1 - cosine similarity.\n  return 1. - cosine_similarity(predictions, targets, epsilon)\n", "path": "optax/_src/loss.py"}], "after_files": [{"content": "# Copyright 2019 DeepMind Technologies Limited. 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.\n# ==============================================================================\n\"\"\"Standard losses used in optimisation.\n\nWe provide implementations of the most canonical losses used in deep\nlearning. These operate transparently on batches, and do not perform any\nreduction over the batch dimensions, leaving it to the user to, for instance,\nmean or sum losses across batch dimensions.\n\"\"\"\n\nfrom typing import Optional\n\nimport chex\nimport jax\nimport jax.numpy as jnp\n\nfrom optax._src import utils\n\n\ndef l2_loss(\n    predictions: chex.Array,\n    targets: Optional[chex.Array] = None,\n) -> chex.Array:\n  \"\"\"Calculates the L2 loss for a set of predictions.\n\n  Note: the 0.5 term is standard in \"Pattern Recognition and Machine Learning\"\n  by Bishop, but not \"The Elements of Statistical Learning\" by Tibshirani.\n\n  References:\n    [Chris Bishop, 2006](https://bit.ly/3eeP0ga)\n\n  Args:\n    predictions: a vector of arbitrary shape.\n    targets: a vector of shape compatible with predictions; if not provides\n      then it is assumed to be zero.\n\n  Returns:\n    the squared error loss.\n  \"\"\"\n  chex.assert_type([predictions], float)\n  errors = (predictions - targets) if (targets is not None) else predictions\n  return 0.5 * (errors)**2\n\n\ndef huber_loss(\n    predictions: chex.Array,\n    targets: Optional[chex.Array] = None,\n    delta: float = 1.) -> chex.Array:\n  \"\"\"Huber loss, similar to L2 loss close to zero, L1 loss away from zero.\n\n  If gradient descent is applied to the `huber loss`, it is equivalent to\n  clipping gradients of an `l2_loss` to `[-delta, delta]` in the backward pass.\n\n  References:\n    [Huber, 1964](www.projecteuclid.org/download/pdf_1/euclid.aoms/1177703732)\n\n  Args:\n    predictions: a vector of arbitrary shape.\n    targets: a vector of shape compatible with predictions; if not provides\n      then it is assumed to be zero.\n    delta: the bounds for the huber loss transformation, defaults at 1.\n\n  Returns:\n    a vector of same shape of `x`.\n  \"\"\"\n  chex.assert_type([predictions], float)\n  errors = (predictions - targets) if (targets is not None) else predictions\n  # 0.5 * err^2                  if |err| <= d\n  # 0.5 * d^2 + d * (|err| - d)  if |err| > d\n  abs_errors = jnp.abs(errors)\n  quadratic = jnp.minimum(abs_errors, delta)\n  # Same as max(abs_x - delta, 0) but avoids potentially doubling gradient.\n  linear = abs_errors - quadratic\n  return 0.5 * quadratic ** 2 + delta * linear\n\n\ndef smooth_labels(\n    labels: chex.Array,\n    alpha: float,\n) -> jnp.ndarray:\n  \"\"\"Apply label smoothing.\n\n  Label smoothing is often used in combination with a cross-entropy loss.\n  Smoothed labels favour small logit gaps, and it has been shown that this can\n  provide better model calibration by preventing overconfident predictions.\n\n  References:\n    [M\u00fcller et al, 2019](https://arxiv.org/pdf/1906.02629.pdf)\n\n  Args:\n    labels: one hot labels to be smoothed.\n    alpha: the smoothing factor, the greedy category with be assigned\n      probability `(1-alpha) + alpha / num_categories`\n\n  Returns:\n    a smoothed version of the one hot input labels.\n\n  \"\"\"\n  chex.assert_type([labels], float)\n  num_categories = labels.shape[-1]\n  return (1.0 - alpha) * labels + alpha / num_categories\n\n\ndef sigmoid_binary_cross_entropy(logits, labels):\n  \"\"\"Computes sigmoid cross entropy given logits and multiple class labels.\n\n  Measures the probability error in discrete classification tasks in which\n  each class is an independent binary prediction and different classes are\n  not mutually exclusive. This may be used for multilabel image classification\n  for instance a model may predict that an image contains both a cat and a dog.\n\n  References:\n    [Goodfellow et al, 2016](http://www.deeplearningbook.org/contents/prob.html)\n\n  Args:\n    logits: unnormalized log probabilities.\n    labels: the probability for that class.\n\n  Returns:\n    a sigmoid cross entropy loss.\n  \"\"\"\n  chex.assert_equal_shape([logits, labels])\n  chex.assert_type([logits, labels], float)\n  log_p = jax.nn.log_sigmoid(logits)\n  # log(1 - sigmoid(x)) = log_sigmoid(-x), the latter more numerically stable\n  log_not_p = jax.nn.log_sigmoid(-logits)\n  return -labels * log_p - (1. - labels) * log_not_p\n\n\ndef softmax_cross_entropy(\n    logits: chex.Array,\n    labels: chex.Array,\n) -> chex.Array:\n  \"\"\"Computes the softmax cross entropy between sets of logits and labels.\n\n  Measures the probability error in discrete classification tasks in which\n  the classes are mutually exclusive (each entry is in exactly one class).\n  For example, each CIFAR-10 image is labeled with one and only one label:\n  an image can be a dog or a truck, but not both.\n\n  References:\n    [Goodfellow et al, 2016](http://www.deeplearningbook.org/contents/prob.html)\n\n  Args:\n    logits: unnormalized log probabilities.\n    labels: a valid probability distribution (non-negative, sum to 1), e.g a\n      one hot encoding of which class is the correct one for each input.\n\n  Returns:\n    the cross entropy loss.\n  \"\"\"\n  chex.assert_equal_shape([logits, labels])\n  chex.assert_type([logits, labels], float)\n  return -jnp.sum(labels * jax.nn.log_softmax(logits, axis=-1), axis=-1)\n\n\ndef cosine_similarity(\n    predictions: chex.Array,\n    targets: chex.Array,\n    epsilon: float = 0.,\n) -> chex.Array:\n  r\"\"\"Computes the cosine similarity between targets and predictions.\n\n  The cosine **similarity** is a measure of similarity between vectors defined\n  as the cosine of the angle between them, which is also the inner product of\n  those vectors normalized to have unit norm.\n\n  References:\n    [Wikipedia, 2021](https://en.wikipedia.org/wiki/Cosine_similarity)\n\n  Args:\n    predictions: The predicted vector.\n    targets: Ground truth target vector.\n    epsilon: minimum norm for terms in the denominator of the cosine similarity.\n\n  Returns:\n    cosine similarity values.\n  \"\"\"\n  chex.assert_equal_shape([targets, predictions])\n  chex.assert_type([targets, predictions], float)\n  # vectorize norm fn, to treat all dimensions except the last as batch dims.\n  batched_norm_fn = jnp.vectorize(\n      utils.safe_norm, signature='(k)->()', excluded={1})\n  # normalise the last dimension of targets and predictions.\n  unit_targets = targets / jnp.expand_dims(\n      batched_norm_fn(targets, epsilon), axis=-1)\n  unit_predictions = predictions / jnp.expand_dims(\n      batched_norm_fn(predictions, epsilon), axis=-1)\n  # return cosine similarity.\n  return jnp.sum(unit_targets * unit_predictions, axis=-1)\n\n\ndef cosine_distance(\n    predictions: chex.Array,\n    targets: chex.Array,\n    epsilon: float = 0.,\n) -> chex.Array:\n  r\"\"\"Computes the cosine distance between targets and predictions.\n\n  The cosine **distance**, implemented here, measures the **dissimilarity**\n  of two vectors as the opposite of cosine **similarity**: `1 - cos(\\theta)`.\n\n  References:\n    [Wikipedia, 2021](https://en.wikipedia.org/wiki/Cosine_similarity)\n\n  Args:\n    predictions: The predicted vector.\n    targets: Ground truth target vector.\n    epsilon: minimum norm for terms in the denominator of the cosine similarity.\n\n  Returns:\n    cosine similarity values.\n  \"\"\"\n  chex.assert_equal_shape([targets, predictions])\n  chex.assert_type([targets, predictions], float)\n  # cosine distance = 1 - cosine similarity.\n  return 1. - cosine_similarity(predictions, targets, epsilon)\n\n\ndef log_cosh(\n    predictions: chex.Array,\n    targets: Optional[chex.Array] = None,\n) -> chex.Array:\n  \"\"\"Calculates the log-cosh loss for a set of predictions.\n\n  log(cosh(x)) is approximately `(x**2) / 2` for small x and `abs(x) - log(2)`\n  for large x.  It is a twice differentiable alternative to the Huber loss.\n\n  References:\n    [Chen et al, 2019](https://openreview.net/pdf?id=rkglvsC9Ym)\n\n  Args:\n    predictions: a vector of arbitrary shape.\n    targets: a vector of shape compatible with predictions; if not provided\n      then it is assumed to be zero.\n\n  Returns:\n    the log-cosh loss.\n  \"\"\"\n  chex.assert_type([predictions], float)\n  errors = (predictions - targets) if (targets is not None) else predictions\n  # log(cosh(x)) = log((exp(x) + exp(-x))/2) = log(exp(x) + exp(-x)) - log(2)\n  return jnp.logaddexp(errors, -errors) - jnp.log(2.0).astype(errors.dtype)\n", "path": "optax/_src/loss.py"}]}