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