We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
`slogdet` sign is incorrect
In the following, numpy disagrees with JAX
```python
mat = np.array([[[-0.01]],
             [[-0.01]]])
print(np.linalg.slogdet(mat))
print(jnp.linalg.slogdet(jnp.array(mat)))
```
=>
```
(array([-1., -1.]), array([-4.60517019, -4.60517019]))
(DeviceArray([1., 1.]), DeviceArray([-4.60517019, -4.60517019]))
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `jax/numpy/linalg.py`
Content:
```
1 # Copyright 2018 Google LLC
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 #     https://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 from __future__ import absolute_import
16 from __future__ import division
17 from __future__ import print_function
18 
19 from functools import partial
20 
21 import numpy as onp
22 import warnings
23 
24 from jax import jit
25 from .. import lax
26 from .. import lax_linalg
27 from .lax_numpy import _not_implemented
28 from .lax_numpy import _wraps
29 from . import lax_numpy as np
30 from ..util import get_module_functions
31 from ..lib import xla_bridge
32 
33 
34 _T = lambda x: np.swapaxes(x, -1, -2)
35 
36 
37 def _promote_arg_dtypes(*args):
38   """Promotes `args` to a common inexact type."""
39   def _to_inexact_type(type):
40     return type if np.issubdtype(type, np.inexact) else np.float64
41   inexact_types = [_to_inexact_type(np._dtype(arg)) for arg in args]
42   dtype = xla_bridge.canonicalize_dtype(np.result_type(*inexact_types))
43   args = [lax.convert_element_type(arg, dtype) for arg in args]
44   if len(args) == 1:
45     return args[0]
46   else:
47     return args
48 
49 
50 @_wraps(onp.linalg.cholesky)
51 def cholesky(a):
52   a = _promote_arg_dtypes(np.asarray(a))
53   return lax_linalg.cholesky(a)
54 
55 
56 @_wraps(onp.linalg.svd)
57 def svd(a, full_matrices=True, compute_uv=True):
58   a = _promote_arg_dtypes(np.asarray(a))
59   return lax_linalg.svd(a, full_matrices, compute_uv)
60 
61 
62 @_wraps(onp.linalg.slogdet)
63 def slogdet(a):
64   a = _promote_arg_dtypes(np.asarray(a))
65   dtype = lax.dtype(a)
66   a_shape = np.shape(a)
67   if len(a_shape) < 2 or a_shape[-1] != a_shape[-2]:
68     msg = "Argument to slogdet() must have shape [..., n, n], got {}"
69     raise ValueError(msg.format(a_shape))
70   lu, pivot = lax_linalg.lu(a)
71   diag = np.diagonal(lu, axis1=-2, axis2=-1)
72   is_zero = np.any(diag == np.array(0, dtype=dtype), axis=-1)
73   parity = np.count_nonzero(pivot != np.arange(a_shape[-1]), axis=-1)
74   if np.iscomplexobj(a):
75     sign = np.prod(diag / np.abs(diag))
76   else:
77     sign = np.array(1, dtype=dtype)
78     parity = parity + np.count_nonzero(diag < 0)
79   sign = np.where(is_zero,
80                   np.array(0, dtype=dtype),
81                   sign * np.array(-2 * (parity % 2) + 1, dtype=dtype))
82   logdet = np.where(
83       is_zero, np.array(-np.inf, dtype=dtype),
84       np.sum(np.log(np.abs(diag)), axis=-1))
85   return sign, np.real(logdet)
86 
87 
88 @_wraps(onp.linalg.det)
89 def det(a):
90   sign, logdet = slogdet(a)
91   return sign * np.exp(logdet)
92 
93 
94 @_wraps(onp.linalg.eig)
95 def eig(a):
96   a = _promote_arg_dtypes(np.asarray(a))
97   w, vl, vr = lax_linalg.eig(a)
98   return w, vr
99 
100 
101 @_wraps(onp.linalg.eigh)
102 def eigh(a, UPLO=None, symmetrize_input=True):
103   if UPLO is None or UPLO == "L":
104     lower = True
105   elif UPLO == "U":
106     lower = False
107   else:
108     msg = "UPLO must be one of None, 'L', or 'U', got {}".format(UPLO)
109     raise ValueError(msg)
110 
111   a = _promote_arg_dtypes(np.asarray(a))
112   v, w = lax_linalg.eigh(a, lower=lower, symmetrize_input=symmetrize_input)
113   return w, v
114 
115 
116 @_wraps(onp.linalg.inv)
117 def inv(a):
118   if np.ndim(a) < 2 or a.shape[-1] != a.shape[-2]:
119     raise ValueError("Argument to inv must have shape [..., n, n], got {}."
120       .format(np.shape(a)))
121   return solve(
122     a, lax.broadcast(np.eye(a.shape[-1], dtype=lax.dtype(a)), a.shape[:-2]))
123 
124 
125 @partial(jit, static_argnums=(1, 2, 3))
126 def _norm(x, ord, axis, keepdims):
127   x = _promote_arg_dtypes(np.asarray(x))
128   x_shape = np.shape(x)
129   ndim = len(x_shape)
130 
131   if axis is None:
132     # NumPy has an undocumented behavior that admits arbitrary rank inputs if
133     # `ord` is None: https://github.com/numpy/numpy/issues/14215
134     if ord is None:
135       return np.sqrt(np.sum(np.real(x * np.conj(x)), keepdims=keepdims))
136     axis = tuple(range(ndim))
137   elif isinstance(axis, tuple):
138     axis = tuple(np._canonicalize_axis(x, ndim) for x in axis)
139   else:
140     axis = (np._canonicalize_axis(axis, ndim),)
141 
142   num_axes = len(axis)
143   if num_axes == 1:
144     if ord is None or ord == 2:
145       return np.sqrt(np.sum(np.real(x * np.conj(x)), axis=axis,
146                             keepdims=keepdims))
147     elif ord == np.inf:
148       return np.amax(np.abs(x), axis=axis, keepdims=keepdims)
149     elif ord == -np.inf:
150       return np.amin(np.abs(x), axis=axis, keepdims=keepdims)
151     elif ord == 0:
152       return np.sum(x != 0, dtype=np.finfo(lax.dtype(x)).dtype,
153                     axis=axis, keepdims=keepdims)
154     elif ord == 1:
155       # Numpy has a special case for ord == 1 as an optimization. We don't
156       # really need the optimization (XLA could do it for us), but the Numpy
157       # code has slightly different type promotion semantics, so we need a
158       # special case too.
159       return np.sum(np.abs(x), axis=axis, keepdims=keepdims)
160     else:
161       return np.power(np.sum(np.abs(x) ** ord, axis=axis, keepdims=keepdims),
162                       1. / ord)
163 
164   elif num_axes == 2:
165     row_axis, col_axis = axis
166     if ord is None or ord in ('f', 'fro'):
167       return np.sqrt(np.sum(np.real(x * np.conj(x)), axis=axis,
168                             keepdims=keepdims))
169     elif ord == 1:
170       if not keepdims and col_axis > row_axis:
171         col_axis -= 1
172       return np.amax(np.sum(np.abs(x), axis=row_axis, keepdims=keepdims),
173                      axis=col_axis, keepdims=keepdims)
174     elif ord == -1:
175       if not keepdims and col_axis > row_axis:
176         col_axis -= 1
177       return np.amin(np.sum(np.abs(x), axis=row_axis, keepdims=keepdims),
178                      axis=col_axis, keepdims=keepdims)
179     elif ord == np.inf:
180       if not keepdims and row_axis > col_axis:
181         row_axis -= 1
182       return np.amax(np.sum(np.abs(x), axis=col_axis, keepdims=keepdims),
183                      axis=row_axis, keepdims=keepdims)
184     elif ord == -np.inf:
185       if not keepdims and row_axis > col_axis:
186         row_axis -= 1
187       return np.amin(np.sum(np.abs(x), axis=col_axis, keepdims=keepdims),
188                      axis=row_axis, keepdims=keepdims)
189     elif ord in ('nuc', 2, -2):
190       x = np.moveaxis(x, axis, (-2, -1))
191       if ord == 2:
192         reducer = np.amax
193       elif ord == -2:
194         reducer = np.amin
195       else:
196         reducer = np.sum
197       y = reducer(svd(x, compute_uv=False), axis=-1)
198       if keepdims:
199         result_shape = list(x_shape)
200         result_shape[axis[0]] = 1
201         result_shape[axis[1]] = 1
202         y = np.reshape(y, result_shape)
203       return y
204     else:
205       raise ValueError("Invalid order '{}' for matrix norm.".format(ord))
206   else:
207     raise ValueError(
208         "Invalid axis values ({}) for np.linalg.norm.".format(axis))
209 
210 @_wraps(onp.linalg.norm)
211 def norm(x, ord=None, axis=None, keepdims=False):
212   return _norm(x, ord, axis, keepdims)
213 
214 
215 @_wraps(onp.linalg.qr)
216 def qr(a, mode="reduced"):
217   if mode in ("reduced", "r", "full"):
218     full_matrices = False
219   elif mode == "complete":
220     full_matrices = True
221   else:
222     raise ValueError("Unsupported QR decomposition mode '{}'".format(mode))
223   a = _promote_arg_dtypes(np.asarray(a))
224   q, r = lax_linalg.qr(a, full_matrices)
225   if mode == "r":
226     return r
227   return q, r
228 
229 
230 @_wraps(onp.linalg.solve)
231 @jit
232 def solve(a, b):
233   a, b = _promote_arg_dtypes(np.asarray(a), np.asarray(b))
234   a_shape = np.shape(a)
235   b_shape = np.shape(b)
236   a_ndims = len(a_shape)
237   b_ndims = len(b_shape)
238   if not (a_ndims >= 2 and a_shape[-1] == a_shape[-2] and b_ndims >= 1):
239     msg = ("The arguments to solve must have shapes a=[..., m, m] and "
240            "b=[..., m, k] or b=[..., m]; got a={} and b={}")
241     raise ValueError(msg.format(a_shape, b_shape))
242   lu, pivots = lax_linalg.lu(a)
243   dtype = lax.dtype(a)
244 
245   m = a_shape[-1]
246 
247   # Numpy treats the RHS as a (batched) vector if the number of dimensions
248   # differ by 1. Otherwise, broadcasting rules apply.
249   x = b[..., None] if a_ndims == b_ndims + 1 else b
250 
251   batch_dims = lax.broadcast_shapes(lu.shape[:-2], x.shape[:-2])
252   x = np.broadcast_to(x, batch_dims + x.shape[-2:])
253   lu = np.broadcast_to(lu, batch_dims + lu.shape[-2:])
254 
255   permutation = lax_linalg.lu_pivots_to_permutation(pivots, m)
256   permutation = np.broadcast_to(permutation, batch_dims + (m,))
257   iotas = np.ix_(*(lax.iota(np.int32, b) for b in batch_dims + (1,)))
258   x = x[iotas[:-1] + (permutation, slice(None))]
259 
260   x = lax_linalg.triangular_solve(lu, x, left_side=True, lower=True,
261                                   unit_diagonal=True)
262   x = lax_linalg.triangular_solve(lu, x, left_side=True, lower=False)
263 
264   return x[..., 0] if a_ndims == b_ndims + 1 else x
265 
266 
267 for func in get_module_functions(onp.linalg):
268   if func.__name__ not in globals():
269     globals()[func.__name__] = _not_implemented(func)
270
```
--- 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/jax/numpy/linalg.py b/jax/numpy/linalg.py
--- a/jax/numpy/linalg.py
+++ b/jax/numpy/linalg.py
@@ -60,6 +60,7 @@
 
 
 @_wraps(onp.linalg.slogdet)
+@jit
 def slogdet(a):
   a = _promote_arg_dtypes(np.asarray(a))
   dtype = lax.dtype(a)
@@ -72,10 +73,10 @@
   is_zero = np.any(diag == np.array(0, dtype=dtype), axis=-1)
   parity = np.count_nonzero(pivot != np.arange(a_shape[-1]), axis=-1)
   if np.iscomplexobj(a):
-    sign = np.prod(diag / np.abs(diag))
+    sign = np.prod(diag / np.abs(diag), axis=-1)
   else:
     sign = np.array(1, dtype=dtype)
-    parity = parity + np.count_nonzero(diag < 0)
+    parity = parity + np.count_nonzero(diag < 0, axis=-1)
   sign = np.where(is_zero,
                   np.array(0, dtype=dtype),
                   sign * np.array(-2 * (parity % 2) + 1, dtype=dtype))

{"golden_diff": "diff --git a/jax/numpy/linalg.py b/jax/numpy/linalg.py\n--- a/jax/numpy/linalg.py\n+++ b/jax/numpy/linalg.py\n@@ -60,6 +60,7 @@\n \n \n @_wraps(onp.linalg.slogdet)\n+@jit\n def slogdet(a):\n   a = _promote_arg_dtypes(np.asarray(a))\n   dtype = lax.dtype(a)\n@@ -72,10 +73,10 @@\n   is_zero = np.any(diag == np.array(0, dtype=dtype), axis=-1)\n   parity = np.count_nonzero(pivot != np.arange(a_shape[-1]), axis=-1)\n   if np.iscomplexobj(a):\n-    sign = np.prod(diag / np.abs(diag))\n+    sign = np.prod(diag / np.abs(diag), axis=-1)\n   else:\n     sign = np.array(1, dtype=dtype)\n-    parity = parity + np.count_nonzero(diag < 0)\n+    parity = parity + np.count_nonzero(diag < 0, axis=-1)\n   sign = np.where(is_zero,\n                   np.array(0, dtype=dtype),\n                   sign * np.array(-2 * (parity % 2) + 1, dtype=dtype))\n", "issue": "`slogdet` sign is incorrect\nIn the following, numpy disagrees with JAX\r\n```python\r\nmat = np.array([[[-0.01]],\r\n             [[-0.01]]])\r\nprint(np.linalg.slogdet(mat))\r\nprint(jnp.linalg.slogdet(jnp.array(mat)))\r\n```\r\n=>\r\n```\r\n(array([-1., -1.]), array([-4.60517019, -4.60517019]))\r\n(DeviceArray([1., 1.]), DeviceArray([-4.60517019, -4.60517019]))\r\n```\n", "before_files": [{"content": "# Copyright 2018 Google LLC\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#     https://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\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom functools import partial\n\nimport numpy as onp\nimport warnings\n\nfrom jax import jit\nfrom .. import lax\nfrom .. import lax_linalg\nfrom .lax_numpy import _not_implemented\nfrom .lax_numpy import _wraps\nfrom . import lax_numpy as np\nfrom ..util import get_module_functions\nfrom ..lib import xla_bridge\n\n\n_T = lambda x: np.swapaxes(x, -1, -2)\n\n\ndef _promote_arg_dtypes(*args):\n  \"\"\"Promotes `args` to a common inexact type.\"\"\"\n  def _to_inexact_type(type):\n    return type if np.issubdtype(type, np.inexact) else np.float64\n  inexact_types = [_to_inexact_type(np._dtype(arg)) for arg in args]\n  dtype = xla_bridge.canonicalize_dtype(np.result_type(*inexact_types))\n  args = [lax.convert_element_type(arg, dtype) for arg in args]\n  if len(args) == 1:\n    return args[0]\n  else:\n    return args\n\n\n@_wraps(onp.linalg.cholesky)\ndef cholesky(a):\n  a = _promote_arg_dtypes(np.asarray(a))\n  return lax_linalg.cholesky(a)\n\n\n@_wraps(onp.linalg.svd)\ndef svd(a, full_matrices=True, compute_uv=True):\n  a = _promote_arg_dtypes(np.asarray(a))\n  return lax_linalg.svd(a, full_matrices, compute_uv)\n\n\n@_wraps(onp.linalg.slogdet)\ndef slogdet(a):\n  a = _promote_arg_dtypes(np.asarray(a))\n  dtype = lax.dtype(a)\n  a_shape = np.shape(a)\n  if len(a_shape) < 2 or a_shape[-1] != a_shape[-2]:\n    msg = \"Argument to slogdet() must have shape [..., n, n], got {}\"\n    raise ValueError(msg.format(a_shape))\n  lu, pivot = lax_linalg.lu(a)\n  diag = np.diagonal(lu, axis1=-2, axis2=-1)\n  is_zero = np.any(diag == np.array(0, dtype=dtype), axis=-1)\n  parity = np.count_nonzero(pivot != np.arange(a_shape[-1]), axis=-1)\n  if np.iscomplexobj(a):\n    sign = np.prod(diag / np.abs(diag))\n  else:\n    sign = np.array(1, dtype=dtype)\n    parity = parity + np.count_nonzero(diag < 0)\n  sign = np.where(is_zero,\n                  np.array(0, dtype=dtype),\n                  sign * np.array(-2 * (parity % 2) + 1, dtype=dtype))\n  logdet = np.where(\n      is_zero, np.array(-np.inf, dtype=dtype),\n      np.sum(np.log(np.abs(diag)), axis=-1))\n  return sign, np.real(logdet)\n\n\n@_wraps(onp.linalg.det)\ndef det(a):\n  sign, logdet = slogdet(a)\n  return sign * np.exp(logdet)\n\n\n@_wraps(onp.linalg.eig)\ndef eig(a):\n  a = _promote_arg_dtypes(np.asarray(a))\n  w, vl, vr = lax_linalg.eig(a)\n  return w, vr\n\n\n@_wraps(onp.linalg.eigh)\ndef eigh(a, UPLO=None, symmetrize_input=True):\n  if UPLO is None or UPLO == \"L\":\n    lower = True\n  elif UPLO == \"U\":\n    lower = False\n  else:\n    msg = \"UPLO must be one of None, 'L', or 'U', got {}\".format(UPLO)\n    raise ValueError(msg)\n\n  a = _promote_arg_dtypes(np.asarray(a))\n  v, w = lax_linalg.eigh(a, lower=lower, symmetrize_input=symmetrize_input)\n  return w, v\n\n\n@_wraps(onp.linalg.inv)\ndef inv(a):\n  if np.ndim(a) < 2 or a.shape[-1] != a.shape[-2]:\n    raise ValueError(\"Argument to inv must have shape [..., n, n], got {}.\"\n      .format(np.shape(a)))\n  return solve(\n    a, lax.broadcast(np.eye(a.shape[-1], dtype=lax.dtype(a)), a.shape[:-2]))\n\n\n@partial(jit, static_argnums=(1, 2, 3))\ndef _norm(x, ord, axis, keepdims):\n  x = _promote_arg_dtypes(np.asarray(x))\n  x_shape = np.shape(x)\n  ndim = len(x_shape)\n\n  if axis is None:\n    # NumPy has an undocumented behavior that admits arbitrary rank inputs if\n    # `ord` is None: https://github.com/numpy/numpy/issues/14215\n    if ord is None:\n      return np.sqrt(np.sum(np.real(x * np.conj(x)), keepdims=keepdims))\n    axis = tuple(range(ndim))\n  elif isinstance(axis, tuple):\n    axis = tuple(np._canonicalize_axis(x, ndim) for x in axis)\n  else:\n    axis = (np._canonicalize_axis(axis, ndim),)\n\n  num_axes = len(axis)\n  if num_axes == 1:\n    if ord is None or ord == 2:\n      return np.sqrt(np.sum(np.real(x * np.conj(x)), axis=axis,\n                            keepdims=keepdims))\n    elif ord == np.inf:\n      return np.amax(np.abs(x), axis=axis, keepdims=keepdims)\n    elif ord == -np.inf:\n      return np.amin(np.abs(x), axis=axis, keepdims=keepdims)\n    elif ord == 0:\n      return np.sum(x != 0, dtype=np.finfo(lax.dtype(x)).dtype,\n                    axis=axis, keepdims=keepdims)\n    elif ord == 1:\n      # Numpy has a special case for ord == 1 as an optimization. We don't\n      # really need the optimization (XLA could do it for us), but the Numpy\n      # code has slightly different type promotion semantics, so we need a\n      # special case too.\n      return np.sum(np.abs(x), axis=axis, keepdims=keepdims)\n    else:\n      return np.power(np.sum(np.abs(x) ** ord, axis=axis, keepdims=keepdims),\n                      1. / ord)\n\n  elif num_axes == 2:\n    row_axis, col_axis = axis\n    if ord is None or ord in ('f', 'fro'):\n      return np.sqrt(np.sum(np.real(x * np.conj(x)), axis=axis,\n                            keepdims=keepdims))\n    elif ord == 1:\n      if not keepdims and col_axis > row_axis:\n        col_axis -= 1\n      return np.amax(np.sum(np.abs(x), axis=row_axis, keepdims=keepdims),\n                     axis=col_axis, keepdims=keepdims)\n    elif ord == -1:\n      if not keepdims and col_axis > row_axis:\n        col_axis -= 1\n      return np.amin(np.sum(np.abs(x), axis=row_axis, keepdims=keepdims),\n                     axis=col_axis, keepdims=keepdims)\n    elif ord == np.inf:\n      if not keepdims and row_axis > col_axis:\n        row_axis -= 1\n      return np.amax(np.sum(np.abs(x), axis=col_axis, keepdims=keepdims),\n                     axis=row_axis, keepdims=keepdims)\n    elif ord == -np.inf:\n      if not keepdims and row_axis > col_axis:\n        row_axis -= 1\n      return np.amin(np.sum(np.abs(x), axis=col_axis, keepdims=keepdims),\n                     axis=row_axis, keepdims=keepdims)\n    elif ord in ('nuc', 2, -2):\n      x = np.moveaxis(x, axis, (-2, -1))\n      if ord == 2:\n        reducer = np.amax\n      elif ord == -2:\n        reducer = np.amin\n      else:\n        reducer = np.sum\n      y = reducer(svd(x, compute_uv=False), axis=-1)\n      if keepdims:\n        result_shape = list(x_shape)\n        result_shape[axis[0]] = 1\n        result_shape[axis[1]] = 1\n        y = np.reshape(y, result_shape)\n      return y\n    else:\n      raise ValueError(\"Invalid order '{}' for matrix norm.\".format(ord))\n  else:\n    raise ValueError(\n        \"Invalid axis values ({}) for np.linalg.norm.\".format(axis))\n\n@_wraps(onp.linalg.norm)\ndef norm(x, ord=None, axis=None, keepdims=False):\n  return _norm(x, ord, axis, keepdims)\n\n\n@_wraps(onp.linalg.qr)\ndef qr(a, mode=\"reduced\"):\n  if mode in (\"reduced\", \"r\", \"full\"):\n    full_matrices = False\n  elif mode == \"complete\":\n    full_matrices = True\n  else:\n    raise ValueError(\"Unsupported QR decomposition mode '{}'\".format(mode))\n  a = _promote_arg_dtypes(np.asarray(a))\n  q, r = lax_linalg.qr(a, full_matrices)\n  if mode == \"r\":\n    return r\n  return q, r\n\n\n@_wraps(onp.linalg.solve)\n@jit\ndef solve(a, b):\n  a, b = _promote_arg_dtypes(np.asarray(a), np.asarray(b))\n  a_shape = np.shape(a)\n  b_shape = np.shape(b)\n  a_ndims = len(a_shape)\n  b_ndims = len(b_shape)\n  if not (a_ndims >= 2 and a_shape[-1] == a_shape[-2] and b_ndims >= 1):\n    msg = (\"The arguments to solve must have shapes a=[..., m, m] and \"\n           \"b=[..., m, k] or b=[..., m]; got a={} and b={}\")\n    raise ValueError(msg.format(a_shape, b_shape))\n  lu, pivots = lax_linalg.lu(a)\n  dtype = lax.dtype(a)\n\n  m = a_shape[-1]\n\n  # Numpy treats the RHS as a (batched) vector if the number of dimensions\n  # differ by 1. Otherwise, broadcasting rules apply.\n  x = b[..., None] if a_ndims == b_ndims + 1 else b\n\n  batch_dims = lax.broadcast_shapes(lu.shape[:-2], x.shape[:-2])\n  x = np.broadcast_to(x, batch_dims + x.shape[-2:])\n  lu = np.broadcast_to(lu, batch_dims + lu.shape[-2:])\n\n  permutation = lax_linalg.lu_pivots_to_permutation(pivots, m)\n  permutation = np.broadcast_to(permutation, batch_dims + (m,))\n  iotas = np.ix_(*(lax.iota(np.int32, b) for b in batch_dims + (1,)))\n  x = x[iotas[:-1] + (permutation, slice(None))]\n\n  x = lax_linalg.triangular_solve(lu, x, left_side=True, lower=True,\n                                  unit_diagonal=True)\n  x = lax_linalg.triangular_solve(lu, x, left_side=True, lower=False)\n\n  return x[..., 0] if a_ndims == b_ndims + 1 else x\n\n\nfor func in get_module_functions(onp.linalg):\n  if func.__name__ not in globals():\n    globals()[func.__name__] = _not_implemented(func)\n", "path": "jax/numpy/linalg.py"}], "after_files": [{"content": "# Copyright 2018 Google LLC\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#     https://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\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom functools import partial\n\nimport numpy as onp\nimport warnings\n\nfrom jax import jit\nfrom .. import lax\nfrom .. import lax_linalg\nfrom .lax_numpy import _not_implemented\nfrom .lax_numpy import _wraps\nfrom . import lax_numpy as np\nfrom ..util import get_module_functions\nfrom ..lib import xla_bridge\n\n\n_T = lambda x: np.swapaxes(x, -1, -2)\n\n\ndef _promote_arg_dtypes(*args):\n  \"\"\"Promotes `args` to a common inexact type.\"\"\"\n  def _to_inexact_type(type):\n    return type if np.issubdtype(type, np.inexact) else np.float64\n  inexact_types = [_to_inexact_type(np._dtype(arg)) for arg in args]\n  dtype = xla_bridge.canonicalize_dtype(np.result_type(*inexact_types))\n  args = [lax.convert_element_type(arg, dtype) for arg in args]\n  if len(args) == 1:\n    return args[0]\n  else:\n    return args\n\n\n@_wraps(onp.linalg.cholesky)\ndef cholesky(a):\n  a = _promote_arg_dtypes(np.asarray(a))\n  return lax_linalg.cholesky(a)\n\n\n@_wraps(onp.linalg.svd)\ndef svd(a, full_matrices=True, compute_uv=True):\n  a = _promote_arg_dtypes(np.asarray(a))\n  return lax_linalg.svd(a, full_matrices, compute_uv)\n\n\n@_wraps(onp.linalg.slogdet)\n@jit\ndef slogdet(a):\n  a = _promote_arg_dtypes(np.asarray(a))\n  dtype = lax.dtype(a)\n  a_shape = np.shape(a)\n  if len(a_shape) < 2 or a_shape[-1] != a_shape[-2]:\n    msg = \"Argument to slogdet() must have shape [..., n, n], got {}\"\n    raise ValueError(msg.format(a_shape))\n  lu, pivot = lax_linalg.lu(a)\n  diag = np.diagonal(lu, axis1=-2, axis2=-1)\n  is_zero = np.any(diag == np.array(0, dtype=dtype), axis=-1)\n  parity = np.count_nonzero(pivot != np.arange(a_shape[-1]), axis=-1)\n  if np.iscomplexobj(a):\n    sign = np.prod(diag / np.abs(diag), axis=-1)\n  else:\n    sign = np.array(1, dtype=dtype)\n    parity = parity + np.count_nonzero(diag < 0, axis=-1)\n  sign = np.where(is_zero,\n                  np.array(0, dtype=dtype),\n                  sign * np.array(-2 * (parity % 2) + 1, dtype=dtype))\n  logdet = np.where(\n      is_zero, np.array(-np.inf, dtype=dtype),\n      np.sum(np.log(np.abs(diag)), axis=-1))\n  return sign, np.real(logdet)\n\n\n@_wraps(onp.linalg.det)\ndef det(a):\n  sign, logdet = slogdet(a)\n  return sign * np.exp(logdet)\n\n\n@_wraps(onp.linalg.eig)\ndef eig(a):\n  a = _promote_arg_dtypes(np.asarray(a))\n  w, vl, vr = lax_linalg.eig(a)\n  return w, vr\n\n\n@_wraps(onp.linalg.eigh)\ndef eigh(a, UPLO=None, symmetrize_input=True):\n  if UPLO is None or UPLO == \"L\":\n    lower = True\n  elif UPLO == \"U\":\n    lower = False\n  else:\n    msg = \"UPLO must be one of None, 'L', or 'U', got {}\".format(UPLO)\n    raise ValueError(msg)\n\n  a = _promote_arg_dtypes(np.asarray(a))\n  v, w = lax_linalg.eigh(a, lower=lower, symmetrize_input=symmetrize_input)\n  return w, v\n\n\n@_wraps(onp.linalg.inv)\ndef inv(a):\n  if np.ndim(a) < 2 or a.shape[-1] != a.shape[-2]:\n    raise ValueError(\"Argument to inv must have shape [..., n, n], got {}.\"\n      .format(np.shape(a)))\n  return solve(\n    a, lax.broadcast(np.eye(a.shape[-1], dtype=lax.dtype(a)), a.shape[:-2]))\n\n\n@partial(jit, static_argnums=(1, 2, 3))\ndef _norm(x, ord, axis, keepdims):\n  x = _promote_arg_dtypes(np.asarray(x))\n  x_shape = np.shape(x)\n  ndim = len(x_shape)\n\n  if axis is None:\n    # NumPy has an undocumented behavior that admits arbitrary rank inputs if\n    # `ord` is None: https://github.com/numpy/numpy/issues/14215\n    if ord is None:\n      return np.sqrt(np.sum(np.real(x * np.conj(x)), keepdims=keepdims))\n    axis = tuple(range(ndim))\n  elif isinstance(axis, tuple):\n    axis = tuple(np._canonicalize_axis(x, ndim) for x in axis)\n  else:\n    axis = (np._canonicalize_axis(axis, ndim),)\n\n  num_axes = len(axis)\n  if num_axes == 1:\n    if ord is None or ord == 2:\n      return np.sqrt(np.sum(np.real(x * np.conj(x)), axis=axis,\n                            keepdims=keepdims))\n    elif ord == np.inf:\n      return np.amax(np.abs(x), axis=axis, keepdims=keepdims)\n    elif ord == -np.inf:\n      return np.amin(np.abs(x), axis=axis, keepdims=keepdims)\n    elif ord == 0:\n      return np.sum(x != 0, dtype=np.finfo(lax.dtype(x)).dtype,\n                    axis=axis, keepdims=keepdims)\n    elif ord == 1:\n      # Numpy has a special case for ord == 1 as an optimization. We don't\n      # really need the optimization (XLA could do it for us), but the Numpy\n      # code has slightly different type promotion semantics, so we need a\n      # special case too.\n      return np.sum(np.abs(x), axis=axis, keepdims=keepdims)\n    else:\n      return np.power(np.sum(np.abs(x) ** ord, axis=axis, keepdims=keepdims),\n                      1. / ord)\n\n  elif num_axes == 2:\n    row_axis, col_axis = axis\n    if ord is None or ord in ('f', 'fro'):\n      return np.sqrt(np.sum(np.real(x * np.conj(x)), axis=axis,\n                            keepdims=keepdims))\n    elif ord == 1:\n      if not keepdims and col_axis > row_axis:\n        col_axis -= 1\n      return np.amax(np.sum(np.abs(x), axis=row_axis, keepdims=keepdims),\n                     axis=col_axis, keepdims=keepdims)\n    elif ord == -1:\n      if not keepdims and col_axis > row_axis:\n        col_axis -= 1\n      return np.amin(np.sum(np.abs(x), axis=row_axis, keepdims=keepdims),\n                     axis=col_axis, keepdims=keepdims)\n    elif ord == np.inf:\n      if not keepdims and row_axis > col_axis:\n        row_axis -= 1\n      return np.amax(np.sum(np.abs(x), axis=col_axis, keepdims=keepdims),\n                     axis=row_axis, keepdims=keepdims)\n    elif ord == -np.inf:\n      if not keepdims and row_axis > col_axis:\n        row_axis -= 1\n      return np.amin(np.sum(np.abs(x), axis=col_axis, keepdims=keepdims),\n                     axis=row_axis, keepdims=keepdims)\n    elif ord in ('nuc', 2, -2):\n      x = np.moveaxis(x, axis, (-2, -1))\n      if ord == 2:\n        reducer = np.amax\n      elif ord == -2:\n        reducer = np.amin\n      else:\n        reducer = np.sum\n      y = reducer(svd(x, compute_uv=False), axis=-1)\n      if keepdims:\n        result_shape = list(x_shape)\n        result_shape[axis[0]] = 1\n        result_shape[axis[1]] = 1\n        y = np.reshape(y, result_shape)\n      return y\n    else:\n      raise ValueError(\"Invalid order '{}' for matrix norm.\".format(ord))\n  else:\n    raise ValueError(\n        \"Invalid axis values ({}) for np.linalg.norm.\".format(axis))\n\n@_wraps(onp.linalg.norm)\ndef norm(x, ord=None, axis=None, keepdims=False):\n  return _norm(x, ord, axis, keepdims)\n\n\n@_wraps(onp.linalg.qr)\ndef qr(a, mode=\"reduced\"):\n  if mode in (\"reduced\", \"r\", \"full\"):\n    full_matrices = False\n  elif mode == \"complete\":\n    full_matrices = True\n  else:\n    raise ValueError(\"Unsupported QR decomposition mode '{}'\".format(mode))\n  a = _promote_arg_dtypes(np.asarray(a))\n  q, r = lax_linalg.qr(a, full_matrices)\n  if mode == \"r\":\n    return r\n  return q, r\n\n\n@_wraps(onp.linalg.solve)\n@jit\ndef solve(a, b):\n  a, b = _promote_arg_dtypes(np.asarray(a), np.asarray(b))\n  a_shape = np.shape(a)\n  b_shape = np.shape(b)\n  a_ndims = len(a_shape)\n  b_ndims = len(b_shape)\n  if not (a_ndims >= 2 and a_shape[-1] == a_shape[-2] and b_ndims >= 1):\n    msg = (\"The arguments to solve must have shapes a=[..., m, m] and \"\n           \"b=[..., m, k] or b=[..., m]; got a={} and b={}\")\n    raise ValueError(msg.format(a_shape, b_shape))\n  lu, pivots = lax_linalg.lu(a)\n  dtype = lax.dtype(a)\n\n  m = a_shape[-1]\n\n  # Numpy treats the RHS as a (batched) vector if the number of dimensions\n  # differ by 1. Otherwise, broadcasting rules apply.\n  x = b[..., None] if a_ndims == b_ndims + 1 else b\n\n  batch_dims = lax.broadcast_shapes(lu.shape[:-2], x.shape[:-2])\n  x = np.broadcast_to(x, batch_dims + x.shape[-2:])\n  lu = np.broadcast_to(lu, batch_dims + lu.shape[-2:])\n\n  permutation = lax_linalg.lu_pivots_to_permutation(pivots, m)\n  permutation = np.broadcast_to(permutation, batch_dims + (m,))\n  iotas = np.ix_(*(lax.iota(np.int32, b) for b in batch_dims + (1,)))\n  x = x[iotas[:-1] + (permutation, slice(None))]\n\n  x = lax_linalg.triangular_solve(lu, x, left_side=True, lower=True,\n                                  unit_diagonal=True)\n  x = lax_linalg.triangular_solve(lu, x, left_side=True, lower=False)\n\n  return x[..., 0] if a_ndims == b_ndims + 1 else x\n\n\nfor func in get_module_functions(onp.linalg):\n  if func.__name__ not in globals():\n    globals()[func.__name__] = _not_implemented(func)\n", "path": "jax/numpy/linalg.py"}]}