We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Better tests for utils.
Optax tests did not catch a problem with one of the type annotations in #367. This is due to `utils` not having good test coverage. 

I'm marking this as "good first issue". Any tests for `utils` would be very welcome! No need to write tests for all of them at once, PRs with only a single test at a time are very welcome.

--- 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/utils.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 """Utility functions for testing."""
16 
17 from typing import Optional, Tuple, Sequence
18 
19 import chex
20 import jax
21 import jax.numpy as jnp
22 import jax.scipy.stats.norm as multivariate_normal
23 
24 from optax._src import linear_algebra
25 from optax._src import numerics
26 
27 
28 def tile_second_to_last_dim(a: chex.Array) -> chex.Array:
29   ones = jnp.ones_like(a)
30   a = jnp.expand_dims(a, axis=-1)
31   return jnp.expand_dims(ones, axis=-2) * a
32 
33 
34 def canonicalize_dtype(
35     dtype: Optional[chex.ArrayDType]) -> Optional[chex.ArrayDType]:
36   """Canonicalise a dtype, skip if None."""
37   if dtype is not None:
38     return jax.dtypes.canonicalize_dtype(dtype)
39   return dtype
40 
41 
42 def cast_tree(tree: chex.ArrayTree,
43               dtype: Optional[chex.ArrayDType]) -> chex.ArrayTree:
44   """Cast tree to given dtype, skip if None."""
45   if dtype is not None:
46     return jax.tree_util.tree_map(lambda t: t.astype(dtype), tree)
47   else:
48     return tree
49 
50 
51 def set_diags(a: chex.Array, new_diags: chex.Array) -> chex.Array:
52   """Set the diagonals of every DxD matrix in an input of shape NxDxD.
53 
54   Args:
55     a: rank 3, tensor NxDxD.
56     new_diags: NxD matrix, the new diagonals of each DxD matrix.
57 
58   Returns:
59     NxDxD tensor, with the same contents as `a` but with the diagonal
60       changed to `new_diags`.
61   """
62   n, d, d1 = a.shape
63   assert d == d1
64 
65   indices1 = jnp.repeat(jnp.arange(n), d)
66   indices2 = jnp.tile(jnp.arange(d), n)
67   indices3 = indices2
68 
69   # Use numpy array setting
70   a = a.at[indices1, indices2, indices3].set(new_diags.flatten())
71   return a
72 
73 
74 class MultiNormalDiagFromLogScale():
75   """MultiNormalDiag which directly exposes its input parameters."""
76 
77   def __init__(self, loc: chex.Array, log_scale: chex.Array):
78     self._log_scale = log_scale
79     self._scale = jnp.exp(log_scale)
80     self._mean = loc
81     self._param_shape = jax.lax.broadcast_shapes(
82         self._mean.shape, self._scale.shape)
83 
84   def sample(self, shape: Sequence[int],
85              seed: chex.PRNGKey) -> chex.Array:
86     sample_shape = tuple(shape) + self._param_shape
87     return jax.random.normal(
88         seed, shape=sample_shape) * self._scale  + self._mean
89 
90   def log_prob(self, x: chex.Array) -> chex.Array:
91     log_prob = multivariate_normal.logpdf(x, loc=self._mean, scale=self._scale)
92     # Sum over parameter axes.
93     sum_axis = [-(i + 1) for i in range(len(self._param_shape))]
94     return jnp.sum(log_prob, axis=sum_axis)
95 
96   @property
97   def log_scale(self) -> chex.Array:
98     return self._log_scale
99 
100   @property
101   def params(self) -> Sequence[chex.Array]:
102     return [self._mean, self._log_scale]
103 
104 
105 def multi_normal(loc: chex.Array,
106                  log_scale: chex.Array) -> MultiNormalDiagFromLogScale:
107   return MultiNormalDiagFromLogScale(loc=loc, log_scale=log_scale)
108 
109 
110 @jax.custom_vjp
111 def _scale_gradient(inputs: chex.ArrayTree, scale: float) -> chex.ArrayTree:
112   """Internal gradient scaling implementation."""
113   del scale  # Only used for the backward pass defined in _scale_gradient_bwd.
114   return inputs
115 
116 
117 def _scale_gradient_fwd(inputs: chex.ArrayTree,
118                         scale: float) -> Tuple[chex.ArrayTree, float]:
119   return _scale_gradient(inputs, scale), scale
120 
121 
122 def _scale_gradient_bwd(scale: float,
123                         g: chex.ArrayTree) -> Tuple[chex.ArrayTree, None]:
124   return (jax.tree_util.tree_map(lambda g_: g_ * scale, g), None)
125 
126 
127 _scale_gradient.defvjp(_scale_gradient_fwd, _scale_gradient_bwd)
128 
129 
130 def scale_gradient(inputs: chex.ArrayTree, scale: float) -> chex.ArrayTree:
131   """Scales gradients for the backwards pass.
132 
133   Args:
134     inputs: A nested array.
135     scale: The scale factor for the gradient on the backwards pass.
136 
137   Returns:
138     An array of the same structure as `inputs`, with scaled backward gradient.
139   """
140   # Special case scales of 1. and 0. for more efficiency.
141   if scale == 1.:
142     return inputs
143   elif scale == 0.:
144     return jax.lax.stop_gradient(inputs)
145   else:
146     return _scale_gradient(inputs, scale)
147 
148 
149 # TODO(b/183800387): remove legacy aliases.
150 safe_norm = numerics.safe_norm
151 safe_int32_increment = numerics.safe_int32_increment
152 global_norm = linear_algebra.global_norm
153
```
--- 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/utils.py b/optax/_src/utils.py
--- a/optax/_src/utils.py
+++ b/optax/_src/utils.py
@@ -59,8 +59,22 @@
     NxDxD tensor, with the same contents as `a` but with the diagonal
       changed to `new_diags`.
   """
+  a_dim, new_diags_dim = len(a.shape), len(new_diags.shape)
+  if a_dim != 3:
+    raise ValueError(f'Expected `a` to be a 3D tensor, got {a_dim}D instead')
+  if new_diags_dim != 2:
+    raise ValueError(
+      f'Expected `new_diags` to be a 2D array, got {new_diags_dim}D instead')
   n, d, d1 = a.shape
-  assert d == d1
+  n_diags, d_diags = new_diags.shape
+  if d != d1:
+    raise ValueError(
+      f'Shape mismatch: expected `a.shape` to be {(n, d, d)}, '
+      f'got {(n, d, d1)} instead')
+  if d_diags != d or n_diags != n:
+    raise ValueError(
+      f'Shape mismatch: expected `new_diags.shape` to be {(n, d)}, '
+      f'got {(n_diags, d_diags)} instead')
 
   indices1 = jnp.repeat(jnp.arange(n), d)
   indices2 = jnp.tile(jnp.arange(d), n)

{"golden_diff": "diff --git a/optax/_src/utils.py b/optax/_src/utils.py\n--- a/optax/_src/utils.py\n+++ b/optax/_src/utils.py\n@@ -59,8 +59,22 @@\n     NxDxD tensor, with the same contents as `a` but with the diagonal\n       changed to `new_diags`.\n   \"\"\"\n+  a_dim, new_diags_dim = len(a.shape), len(new_diags.shape)\n+  if a_dim != 3:\n+    raise ValueError(f'Expected `a` to be a 3D tensor, got {a_dim}D instead')\n+  if new_diags_dim != 2:\n+    raise ValueError(\n+      f'Expected `new_diags` to be a 2D array, got {new_diags_dim}D instead')\n   n, d, d1 = a.shape\n-  assert d == d1\n+  n_diags, d_diags = new_diags.shape\n+  if d != d1:\n+    raise ValueError(\n+      f'Shape mismatch: expected `a.shape` to be {(n, d, d)}, '\n+      f'got {(n, d, d1)} instead')\n+  if d_diags != d or n_diags != n:\n+    raise ValueError(\n+      f'Shape mismatch: expected `new_diags.shape` to be {(n, d)}, '\n+      f'got {(n_diags, d_diags)} instead')\n \n   indices1 = jnp.repeat(jnp.arange(n), d)\n   indices2 = jnp.tile(jnp.arange(d), n)\n", "issue": "Better tests for utils.\nOptax tests did not catch a problem with one of the type annotations in #367. This is due to `utils` not having good test coverage. \r\n\r\nI'm marking this as \"good first issue\". Any tests for `utils` would be very welcome! No need to write tests for all of them at once, PRs with only a single test at a time are very welcome.\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\"\"\"Utility functions for testing.\"\"\"\n\nfrom typing import Optional, Tuple, Sequence\n\nimport chex\nimport jax\nimport jax.numpy as jnp\nimport jax.scipy.stats.norm as multivariate_normal\n\nfrom optax._src import linear_algebra\nfrom optax._src import numerics\n\n\ndef tile_second_to_last_dim(a: chex.Array) -> chex.Array:\n  ones = jnp.ones_like(a)\n  a = jnp.expand_dims(a, axis=-1)\n  return jnp.expand_dims(ones, axis=-2) * a\n\n\ndef canonicalize_dtype(\n    dtype: Optional[chex.ArrayDType]) -> Optional[chex.ArrayDType]:\n  \"\"\"Canonicalise a dtype, skip if None.\"\"\"\n  if dtype is not None:\n    return jax.dtypes.canonicalize_dtype(dtype)\n  return dtype\n\n\ndef cast_tree(tree: chex.ArrayTree,\n              dtype: Optional[chex.ArrayDType]) -> chex.ArrayTree:\n  \"\"\"Cast tree to given dtype, skip if None.\"\"\"\n  if dtype is not None:\n    return jax.tree_util.tree_map(lambda t: t.astype(dtype), tree)\n  else:\n    return tree\n\n\ndef set_diags(a: chex.Array, new_diags: chex.Array) -> chex.Array:\n  \"\"\"Set the diagonals of every DxD matrix in an input of shape NxDxD.\n\n  Args:\n    a: rank 3, tensor NxDxD.\n    new_diags: NxD matrix, the new diagonals of each DxD matrix.\n\n  Returns:\n    NxDxD tensor, with the same contents as `a` but with the diagonal\n      changed to `new_diags`.\n  \"\"\"\n  n, d, d1 = a.shape\n  assert d == d1\n\n  indices1 = jnp.repeat(jnp.arange(n), d)\n  indices2 = jnp.tile(jnp.arange(d), n)\n  indices3 = indices2\n\n  # Use numpy array setting\n  a = a.at[indices1, indices2, indices3].set(new_diags.flatten())\n  return a\n\n\nclass MultiNormalDiagFromLogScale():\n  \"\"\"MultiNormalDiag which directly exposes its input parameters.\"\"\"\n\n  def __init__(self, loc: chex.Array, log_scale: chex.Array):\n    self._log_scale = log_scale\n    self._scale = jnp.exp(log_scale)\n    self._mean = loc\n    self._param_shape = jax.lax.broadcast_shapes(\n        self._mean.shape, self._scale.shape)\n\n  def sample(self, shape: Sequence[int],\n             seed: chex.PRNGKey) -> chex.Array:\n    sample_shape = tuple(shape) + self._param_shape\n    return jax.random.normal(\n        seed, shape=sample_shape) * self._scale  + self._mean\n\n  def log_prob(self, x: chex.Array) -> chex.Array:\n    log_prob = multivariate_normal.logpdf(x, loc=self._mean, scale=self._scale)\n    # Sum over parameter axes.\n    sum_axis = [-(i + 1) for i in range(len(self._param_shape))]\n    return jnp.sum(log_prob, axis=sum_axis)\n\n  @property\n  def log_scale(self) -> chex.Array:\n    return self._log_scale\n\n  @property\n  def params(self) -> Sequence[chex.Array]:\n    return [self._mean, self._log_scale]\n\n\ndef multi_normal(loc: chex.Array,\n                 log_scale: chex.Array) -> MultiNormalDiagFromLogScale:\n  return MultiNormalDiagFromLogScale(loc=loc, log_scale=log_scale)\n\n\n@jax.custom_vjp\ndef _scale_gradient(inputs: chex.ArrayTree, scale: float) -> chex.ArrayTree:\n  \"\"\"Internal gradient scaling implementation.\"\"\"\n  del scale  # Only used for the backward pass defined in _scale_gradient_bwd.\n  return inputs\n\n\ndef _scale_gradient_fwd(inputs: chex.ArrayTree,\n                        scale: float) -> Tuple[chex.ArrayTree, float]:\n  return _scale_gradient(inputs, scale), scale\n\n\ndef _scale_gradient_bwd(scale: float,\n                        g: chex.ArrayTree) -> Tuple[chex.ArrayTree, None]:\n  return (jax.tree_util.tree_map(lambda g_: g_ * scale, g), None)\n\n\n_scale_gradient.defvjp(_scale_gradient_fwd, _scale_gradient_bwd)\n\n\ndef scale_gradient(inputs: chex.ArrayTree, scale: float) -> chex.ArrayTree:\n  \"\"\"Scales gradients for the backwards pass.\n\n  Args:\n    inputs: A nested array.\n    scale: The scale factor for the gradient on the backwards pass.\n\n  Returns:\n    An array of the same structure as `inputs`, with scaled backward gradient.\n  \"\"\"\n  # Special case scales of 1. and 0. for more efficiency.\n  if scale == 1.:\n    return inputs\n  elif scale == 0.:\n    return jax.lax.stop_gradient(inputs)\n  else:\n    return _scale_gradient(inputs, scale)\n\n\n# TODO(b/183800387): remove legacy aliases.\nsafe_norm = numerics.safe_norm\nsafe_int32_increment = numerics.safe_int32_increment\nglobal_norm = linear_algebra.global_norm\n", "path": "optax/_src/utils.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\"\"\"Utility functions for testing.\"\"\"\n\nfrom typing import Optional, Tuple, Sequence\n\nimport chex\nimport jax\nimport jax.numpy as jnp\nimport jax.scipy.stats.norm as multivariate_normal\n\nfrom optax._src import linear_algebra\nfrom optax._src import numerics\n\n\ndef tile_second_to_last_dim(a: chex.Array) -> chex.Array:\n  ones = jnp.ones_like(a)\n  a = jnp.expand_dims(a, axis=-1)\n  return jnp.expand_dims(ones, axis=-2) * a\n\n\ndef canonicalize_dtype(\n    dtype: Optional[chex.ArrayDType]) -> Optional[chex.ArrayDType]:\n  \"\"\"Canonicalise a dtype, skip if None.\"\"\"\n  if dtype is not None:\n    return jax.dtypes.canonicalize_dtype(dtype)\n  return dtype\n\n\ndef cast_tree(tree: chex.ArrayTree,\n              dtype: Optional[chex.ArrayDType]) -> chex.ArrayTree:\n  \"\"\"Cast tree to given dtype, skip if None.\"\"\"\n  if dtype is not None:\n    return jax.tree_util.tree_map(lambda t: t.astype(dtype), tree)\n  else:\n    return tree\n\n\ndef set_diags(a: chex.Array, new_diags: chex.Array) -> chex.Array:\n  \"\"\"Set the diagonals of every DxD matrix in an input of shape NxDxD.\n\n  Args:\n    a: rank 3, tensor NxDxD.\n    new_diags: NxD matrix, the new diagonals of each DxD matrix.\n\n  Returns:\n    NxDxD tensor, with the same contents as `a` but with the diagonal\n      changed to `new_diags`.\n  \"\"\"\n  a_dim, new_diags_dim = len(a.shape), len(new_diags.shape)\n  if a_dim != 3:\n    raise ValueError(f'Expected `a` to be a 3D tensor, got {a_dim}D instead')\n  if new_diags_dim != 2:\n    raise ValueError(\n      f'Expected `new_diags` to be a 2D array, got {new_diags_dim}D instead')\n  n, d, d1 = a.shape\n  n_diags, d_diags = new_diags.shape\n  if d != d1:\n    raise ValueError(\n      f'Shape mismatch: expected `a.shape` to be {(n, d, d)}, '\n      f'got {(n, d, d1)} instead')\n  if d_diags != d or n_diags != n:\n    raise ValueError(\n      f'Shape mismatch: expected `new_diags.shape` to be {(n, d)}, '\n      f'got {(n_diags, d_diags)} instead')\n\n  indices1 = jnp.repeat(jnp.arange(n), d)\n  indices2 = jnp.tile(jnp.arange(d), n)\n  indices3 = indices2\n\n  # Use numpy array setting\n  a = a.at[indices1, indices2, indices3].set(new_diags.flatten())\n  return a\n\n\nclass MultiNormalDiagFromLogScale():\n  \"\"\"MultiNormalDiag which directly exposes its input parameters.\"\"\"\n\n  def __init__(self, loc: chex.Array, log_scale: chex.Array):\n    self._log_scale = log_scale\n    self._scale = jnp.exp(log_scale)\n    self._mean = loc\n    self._param_shape = jax.lax.broadcast_shapes(\n        self._mean.shape, self._scale.shape)\n\n  def sample(self, shape: Sequence[int],\n             seed: chex.PRNGKey) -> chex.Array:\n    sample_shape = tuple(shape) + self._param_shape\n    return jax.random.normal(\n        seed, shape=sample_shape) * self._scale  + self._mean\n\n  def log_prob(self, x: chex.Array) -> chex.Array:\n    log_prob = multivariate_normal.logpdf(x, loc=self._mean, scale=self._scale)\n    # Sum over parameter axes.\n    sum_axis = [-(i + 1) for i in range(len(self._param_shape))]\n    return jnp.sum(log_prob, axis=sum_axis)\n\n  @property\n  def log_scale(self) -> chex.Array:\n    return self._log_scale\n\n  @property\n  def params(self) -> Sequence[chex.Array]:\n    return [self._mean, self._log_scale]\n\n\ndef multi_normal(loc: chex.Array,\n                 log_scale: chex.Array) -> MultiNormalDiagFromLogScale:\n  return MultiNormalDiagFromLogScale(loc=loc, log_scale=log_scale)\n\n\n@jax.custom_vjp\ndef _scale_gradient(inputs: chex.ArrayTree, scale: float) -> chex.ArrayTree:\n  \"\"\"Internal gradient scaling implementation.\"\"\"\n  del scale  # Only used for the backward pass defined in _scale_gradient_bwd.\n  return inputs\n\n\ndef _scale_gradient_fwd(inputs: chex.ArrayTree,\n                        scale: float) -> Tuple[chex.ArrayTree, float]:\n  return _scale_gradient(inputs, scale), scale\n\n\ndef _scale_gradient_bwd(scale: float,\n                        g: chex.ArrayTree) -> Tuple[chex.ArrayTree, None]:\n  return (jax.tree_util.tree_map(lambda g_: g_ * scale, g), None)\n\n\n_scale_gradient.defvjp(_scale_gradient_fwd, _scale_gradient_bwd)\n\n\ndef scale_gradient(inputs: chex.ArrayTree, scale: float) -> chex.ArrayTree:\n  \"\"\"Scales gradients for the backwards pass.\n\n  Args:\n    inputs: A nested array.\n    scale: The scale factor for the gradient on the backwards pass.\n\n  Returns:\n    An array of the same structure as `inputs`, with scaled backward gradient.\n  \"\"\"\n  # Special case scales of 1. and 0. for more efficiency.\n  if scale == 1.:\n    return inputs\n  elif scale == 0.:\n    return jax.lax.stop_gradient(inputs)\n  else:\n    return _scale_gradient(inputs, scale)\n\n\n# TODO(b/183800387): remove legacy aliases.\nsafe_norm = numerics.safe_norm\nsafe_int32_increment = numerics.safe_int32_increment\nglobal_norm = linear_algebra.global_norm\n", "path": "optax/_src/utils.py"}]}