prompt: 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()
  
  ```
  
golden_diff: 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)
  
verification_info: {"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"}]}