We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
jax missing scipy.special.expit
Would be possible to add gradients for `expit` and `logit`?

--- 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/scipy/special.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 import scipy.special as osp_special
20 
21 from .. import lax
22 from ..numpy.lax_numpy import _wraps
23 
24 
25 # need to create new functions because _wraps sets the __name__ attribute
26 gammaln = _wraps(osp_special.gammaln)(lambda x: lax.lgamma(x))
27 digamma = _wraps(osp_special.digamma)(lambda x: lax.digamma(x))
28 erf = _wraps(osp_special.erf)(lambda x: lax.erf(x))
29 erfc = _wraps(osp_special.erfc)(lambda x: lax.erfc(x))
30 erfinv = _wraps(osp_special.erfinv)(lambda x: lax.erf_inv(x))
31
```
--- 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/scipy/special.py b/jax/scipy/special.py
--- a/jax/scipy/special.py
+++ b/jax/scipy/special.py
@@ -19,7 +19,7 @@
 import scipy.special as osp_special
 
 from .. import lax
-from ..numpy.lax_numpy import _wraps
+from ..numpy.lax_numpy import _wraps, asarray
 
 
 # need to create new functions because _wraps sets the __name__ attribute
@@ -28,3 +28,16 @@
 erf = _wraps(osp_special.erf)(lambda x: lax.erf(x))
 erfc = _wraps(osp_special.erfc)(lambda x: lax.erfc(x))
 erfinv = _wraps(osp_special.erfinv)(lambda x: lax.erf_inv(x))
+
+
+@_wraps(osp_special.logit)
+def logit(x):
+  x = asarray(x)
+  return lax.log(lax.div(x, lax.sub(lax._const(x, 1), x)))
+
+
+@_wraps(osp_special.expit)
+def expit(x):
+  x = asarray(x)
+  one = lax._const(x, 1)
+  return lax.div(one, lax.add(one, lax.exp(lax.neg(x))))

{"golden_diff": "diff --git a/jax/scipy/special.py b/jax/scipy/special.py\n--- a/jax/scipy/special.py\n+++ b/jax/scipy/special.py\n@@ -19,7 +19,7 @@\n import scipy.special as osp_special\n \n from .. import lax\n-from ..numpy.lax_numpy import _wraps\n+from ..numpy.lax_numpy import _wraps, asarray\n \n \n # need to create new functions because _wraps sets the __name__ attribute\n@@ -28,3 +28,16 @@\n erf = _wraps(osp_special.erf)(lambda x: lax.erf(x))\n erfc = _wraps(osp_special.erfc)(lambda x: lax.erfc(x))\n erfinv = _wraps(osp_special.erfinv)(lambda x: lax.erf_inv(x))\n+\n+\n+@_wraps(osp_special.logit)\n+def logit(x):\n+  x = asarray(x)\n+  return lax.log(lax.div(x, lax.sub(lax._const(x, 1), x)))\n+\n+\n+@_wraps(osp_special.expit)\n+def expit(x):\n+  x = asarray(x)\n+  one = lax._const(x, 1)\n+  return lax.div(one, lax.add(one, lax.exp(lax.neg(x))))\n", "issue": "jax missing scipy.special.expit\nWould be possible to add gradients for `expit` and `logit`?\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\nimport scipy.special as osp_special\n\nfrom .. import lax\nfrom ..numpy.lax_numpy import _wraps\n\n\n# need to create new functions because _wraps sets the __name__ attribute\ngammaln = _wraps(osp_special.gammaln)(lambda x: lax.lgamma(x))\ndigamma = _wraps(osp_special.digamma)(lambda x: lax.digamma(x))\nerf = _wraps(osp_special.erf)(lambda x: lax.erf(x))\nerfc = _wraps(osp_special.erfc)(lambda x: lax.erfc(x))\nerfinv = _wraps(osp_special.erfinv)(lambda x: lax.erf_inv(x))\n", "path": "jax/scipy/special.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\nimport scipy.special as osp_special\n\nfrom .. import lax\nfrom ..numpy.lax_numpy import _wraps, asarray\n\n\n# need to create new functions because _wraps sets the __name__ attribute\ngammaln = _wraps(osp_special.gammaln)(lambda x: lax.lgamma(x))\ndigamma = _wraps(osp_special.digamma)(lambda x: lax.digamma(x))\nerf = _wraps(osp_special.erf)(lambda x: lax.erf(x))\nerfc = _wraps(osp_special.erfc)(lambda x: lax.erfc(x))\nerfinv = _wraps(osp_special.erfinv)(lambda x: lax.erf_inv(x))\n\n\n@_wraps(osp_special.logit)\ndef logit(x):\n  x = asarray(x)\n  return lax.log(lax.div(x, lax.sub(lax._const(x, 1), x)))\n\n\n@_wraps(osp_special.expit)\ndef expit(x):\n  x = asarray(x)\n  one = lax._const(x, 1)\n  return lax.div(one, lax.add(one, lax.exp(lax.neg(x))))\n", "path": "jax/scipy/special.py"}]}