prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Perplexity dtype restriction too strict
  ## 🐛 Bug
  
  The perplexity metric requires the input dtype to be either fp32 or fp64, but this doesn't work with e.g. fp16, and users need to manually recast.
  
  https://github.com/Lightning-AI/torchmetrics/blob/894de4caeeae820f60f1871d75334873241e5633/src/torchmetrics/functional/text/perplexity.py#L21
  
  ### Expected behavior
  
  The metric should accept other floating point dtypes.
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `src/torchmetrics/functional/text/perplexity.py`
  Content:
  ```
  1 # Copyright The Lightning team.
  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 from typing import Optional, Tuple
  16 
  17 import torch
  18 from torch import Tensor
  19 from torch.nn import functional as F  # noqa: N812
  20 
  21 _TORCH_FLOAT_OR_DOUBLE = (torch.float32, torch.float64)
  22 
  23 
  24 def _check_shape_and_type_consistency(preds: Tensor, target: Tensor) -> None:
  25     """Check shape and type consistency of input vectors.
  26 
  27     Args:
  28         preds:
  29             Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,
  30             vocab_size]. Scores will be normalized internally using softmax.
  31         target:
  32             Ground truth values with a shape [batch_size, seq_len].
  33 
  34     Raises:
  35         ValueError:
  36             If ``preds`` tensor has no 3 dimensions.
  37         ValueError:
  38             If ``target`` tensor has no 2 dimensions.
  39         ValueError:
  40             If the first two dimensions of ``preds`` and ``target`` do not equal.
  41         TypeError:
  42             If ``preds`` dtype is not one of ``(torch.float16, torch.float32, torch.float64)``
  43         TypeError:
  44             If ``target`` is not of a type LongTensor (torch.int64)
  45 
  46     """
  47     if len(preds.shape) != 3:
  48         raise ValueError(
  49             "Input tensor `preds` is expected to have 3 dimensions, [batch_size, seq_len, vocab_size],"
  50             f" but got {len(preds.shape)}."
  51         )
  52     if len(target.shape) != 2:
  53         raise ValueError(
  54             "Input tensor `target` is expected to have 2 dimensions, [batch_size, seq_len],"
  55             f" but got {len(target.shape)}."
  56         )
  57     if preds.shape[:2] != target.shape:
  58         raise ValueError(
  59             "Input tensors `preds` and `target` are expected to have equaling first two dimensions,"
  60             f" [batch_size, seq_len], but got {preds.shape[:2]} and {target.shape}."
  61         )
  62     if preds.dtype not in _TORCH_FLOAT_OR_DOUBLE:
  63         raise TypeError(
  64             f"Input tensor `preds` is expected to be of a type one of {_TORCH_FLOAT_OR_DOUBLE} but got {preds.dtype}."
  65         )
  66     if target.dtype != torch.int64:
  67         raise TypeError(f"Input tensor `target` is expected to be of a type {torch.int64} but got {target.dtype}.")
  68 
  69 
  70 def _perplexity_update(preds: Tensor, target: Tensor, ignore_index: Optional[int] = None) -> Tuple[Tensor, Tensor]:
  71     """Compute intermediate statistics for Perplexity.
  72 
  73     Args:
  74         preds:
  75             Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,
  76             vocab_size]. Scores will be normalized internally using softmax.
  77         target:
  78             Ground truth values with a shape [batch_size, seq_len].
  79         ignore_index:
  80             Integer specifying a target class to ignore. If given, this class index does not contribute
  81             to the returned score.
  82 
  83     Returns:
  84         Log probabilities, summed over all samples
  85         Number of samples
  86 
  87     """
  88     _check_shape_and_type_consistency(preds, target)
  89 
  90     probs = F.softmax(preds.reshape(-1, preds.shape[-1]), dim=1)
  91     target = target.reshape(-1)
  92 
  93     if ignore_index is not None:
  94         mask = target.ne(ignore_index)
  95         target = target.where(target != ignore_index, torch.tensor(0, device=target.device))
  96     else:
  97         mask = torch.ones_like(target, dtype=torch.bool)
  98 
  99     probs = probs[:, target].diagonal()[mask]
  100     total_log_probs = -probs.log().sum()
  101     count = mask.sum()
  102 
  103     return total_log_probs, count
  104 
  105 
  106 def _perplexity_compute(total: Tensor, count: Tensor) -> Tensor:
  107     """Compute the Perplexity.
  108 
  109     Args:
  110         total: Log probabilities, summed over all samples
  111         count: Number of samples
  112     Returns:
  113         Perplexity
  114 
  115     """
  116     return torch.exp(total / count)
  117 
  118 
  119 def perplexity(preds: Tensor, target: Tensor, ignore_index: Optional[int] = None) -> Tensor:
  120     """Perplexity measures how well a language model predicts a text sample.
  121 
  122     This metric is calculated as the average number of bits per word a model needs to represent the sample.
  123 
  124     Args:
  125         preds:
  126             Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,
  127             vocab_size], which is the output of a language model. Scores will be normalized internally using softmax.
  128         target:
  129             Ground truth values with a shape [batch_size, seq_len].
  130         ignore_index:
  131             Integer specifying a target class to ignore. If given, this class index does not contribute
  132             to the returned score.
  133 
  134     Returns:
  135         Perplexity value
  136 
  137     Examples:
  138         >>> import torch
  139         >>> gen = torch.manual_seed(42)
  140         >>> preds = torch.rand(2, 8, 5, generator=gen)
  141         >>> target = torch.randint(5, (2, 8), generator=gen)
  142         >>> target[0, 6:] = -100
  143         >>> perplexity(preds, target, ignore_index=-100)
  144         tensor(5.8540)
  145 
  146     """
  147     total, count = _perplexity_update(preds, target, ignore_index)
  148     return _perplexity_compute(total, count)
  149
  ```
  --- 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/src/torchmetrics/functional/text/perplexity.py b/src/torchmetrics/functional/text/perplexity.py
  --- a/src/torchmetrics/functional/text/perplexity.py
  +++ b/src/torchmetrics/functional/text/perplexity.py
  @@ -16,9 +16,6 @@
   
   import torch
   from torch import Tensor
  -from torch.nn import functional as F  # noqa: N812
  -
  -_TORCH_FLOAT_OR_DOUBLE = (torch.float32, torch.float64)
   
   
   def _check_shape_and_type_consistency(preds: Tensor, target: Tensor) -> None:
  @@ -59,10 +56,8 @@
               "Input tensors `preds` and `target` are expected to have equaling first two dimensions,"
               f" [batch_size, seq_len], but got {preds.shape[:2]} and {target.shape}."
           )
  -    if preds.dtype not in _TORCH_FLOAT_OR_DOUBLE:
  -        raise TypeError(
  -            f"Input tensor `preds` is expected to be of a type one of {_TORCH_FLOAT_OR_DOUBLE} but got {preds.dtype}."
  -        )
  +    if not preds.is_floating_point():
  +        raise TypeError(f"Input tensor `preds` is expected to be of floating point type but got {preds.dtype}.")
       if target.dtype != torch.int64:
           raise TypeError(f"Input tensor `target` is expected to be of a type {torch.int64} but got {target.dtype}.")
   
  @@ -87,7 +82,7 @@
       """
       _check_shape_and_type_consistency(preds, target)
   
  -    probs = F.softmax(preds.reshape(-1, preds.shape[-1]), dim=1)
  +    probs = torch.nn.functional.softmax(preds.reshape(-1, preds.shape[-1]), dim=1)
       target = target.reshape(-1)
   
       if ignore_index is not None:
  
verification_info: {"golden_diff": "diff --git a/src/torchmetrics/functional/text/perplexity.py b/src/torchmetrics/functional/text/perplexity.py\n--- a/src/torchmetrics/functional/text/perplexity.py\n+++ b/src/torchmetrics/functional/text/perplexity.py\n@@ -16,9 +16,6 @@\n \n import torch\n from torch import Tensor\n-from torch.nn import functional as F  # noqa: N812\n-\n-_TORCH_FLOAT_OR_DOUBLE = (torch.float32, torch.float64)\n \n \n def _check_shape_and_type_consistency(preds: Tensor, target: Tensor) -> None:\n@@ -59,10 +56,8 @@\n             \"Input tensors `preds` and `target` are expected to have equaling first two dimensions,\"\n             f\" [batch_size, seq_len], but got {preds.shape[:2]} and {target.shape}.\"\n         )\n-    if preds.dtype not in _TORCH_FLOAT_OR_DOUBLE:\n-        raise TypeError(\n-            f\"Input tensor `preds` is expected to be of a type one of {_TORCH_FLOAT_OR_DOUBLE} but got {preds.dtype}.\"\n-        )\n+    if not preds.is_floating_point():\n+        raise TypeError(f\"Input tensor `preds` is expected to be of floating point type but got {preds.dtype}.\")\n     if target.dtype != torch.int64:\n         raise TypeError(f\"Input tensor `target` is expected to be of a type {torch.int64} but got {target.dtype}.\")\n \n@@ -87,7 +82,7 @@\n     \"\"\"\n     _check_shape_and_type_consistency(preds, target)\n \n-    probs = F.softmax(preds.reshape(-1, preds.shape[-1]), dim=1)\n+    probs = torch.nn.functional.softmax(preds.reshape(-1, preds.shape[-1]), dim=1)\n     target = target.reshape(-1)\n \n     if ignore_index is not None:\n", "issue": "Perplexity dtype restriction too strict\n## \ud83d\udc1b Bug\r\n\r\nThe perplexity metric requires the input dtype to be either fp32 or fp64, but this doesn't work with e.g. fp16, and users need to manually recast.\r\n\r\nhttps://github.com/Lightning-AI/torchmetrics/blob/894de4caeeae820f60f1871d75334873241e5633/src/torchmetrics/functional/text/perplexity.py#L21\r\n\r\n### Expected behavior\r\n\r\nThe metric should accept other floating point dtypes.\n", "before_files": [{"content": "# Copyright The Lightning team.\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\nfrom typing import Optional, Tuple\n\nimport torch\nfrom torch import Tensor\nfrom torch.nn import functional as F  # noqa: N812\n\n_TORCH_FLOAT_OR_DOUBLE = (torch.float32, torch.float64)\n\n\ndef _check_shape_and_type_consistency(preds: Tensor, target: Tensor) -> None:\n    \"\"\"Check shape and type consistency of input vectors.\n\n    Args:\n        preds:\n            Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,\n            vocab_size]. Scores will be normalized internally using softmax.\n        target:\n            Ground truth values with a shape [batch_size, seq_len].\n\n    Raises:\n        ValueError:\n            If ``preds`` tensor has no 3 dimensions.\n        ValueError:\n            If ``target`` tensor has no 2 dimensions.\n        ValueError:\n            If the first two dimensions of ``preds`` and ``target`` do not equal.\n        TypeError:\n            If ``preds`` dtype is not one of ``(torch.float16, torch.float32, torch.float64)``\n        TypeError:\n            If ``target`` is not of a type LongTensor (torch.int64)\n\n    \"\"\"\n    if len(preds.shape) != 3:\n        raise ValueError(\n            \"Input tensor `preds` is expected to have 3 dimensions, [batch_size, seq_len, vocab_size],\"\n            f\" but got {len(preds.shape)}.\"\n        )\n    if len(target.shape) != 2:\n        raise ValueError(\n            \"Input tensor `target` is expected to have 2 dimensions, [batch_size, seq_len],\"\n            f\" but got {len(target.shape)}.\"\n        )\n    if preds.shape[:2] != target.shape:\n        raise ValueError(\n            \"Input tensors `preds` and `target` are expected to have equaling first two dimensions,\"\n            f\" [batch_size, seq_len], but got {preds.shape[:2]} and {target.shape}.\"\n        )\n    if preds.dtype not in _TORCH_FLOAT_OR_DOUBLE:\n        raise TypeError(\n            f\"Input tensor `preds` is expected to be of a type one of {_TORCH_FLOAT_OR_DOUBLE} but got {preds.dtype}.\"\n        )\n    if target.dtype != torch.int64:\n        raise TypeError(f\"Input tensor `target` is expected to be of a type {torch.int64} but got {target.dtype}.\")\n\n\ndef _perplexity_update(preds: Tensor, target: Tensor, ignore_index: Optional[int] = None) -> Tuple[Tensor, Tensor]:\n    \"\"\"Compute intermediate statistics for Perplexity.\n\n    Args:\n        preds:\n            Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,\n            vocab_size]. Scores will be normalized internally using softmax.\n        target:\n            Ground truth values with a shape [batch_size, seq_len].\n        ignore_index:\n            Integer specifying a target class to ignore. If given, this class index does not contribute\n            to the returned score.\n\n    Returns:\n        Log probabilities, summed over all samples\n        Number of samples\n\n    \"\"\"\n    _check_shape_and_type_consistency(preds, target)\n\n    probs = F.softmax(preds.reshape(-1, preds.shape[-1]), dim=1)\n    target = target.reshape(-1)\n\n    if ignore_index is not None:\n        mask = target.ne(ignore_index)\n        target = target.where(target != ignore_index, torch.tensor(0, device=target.device))\n    else:\n        mask = torch.ones_like(target, dtype=torch.bool)\n\n    probs = probs[:, target].diagonal()[mask]\n    total_log_probs = -probs.log().sum()\n    count = mask.sum()\n\n    return total_log_probs, count\n\n\ndef _perplexity_compute(total: Tensor, count: Tensor) -> Tensor:\n    \"\"\"Compute the Perplexity.\n\n    Args:\n        total: Log probabilities, summed over all samples\n        count: Number of samples\n    Returns:\n        Perplexity\n\n    \"\"\"\n    return torch.exp(total / count)\n\n\ndef perplexity(preds: Tensor, target: Tensor, ignore_index: Optional[int] = None) -> Tensor:\n    \"\"\"Perplexity measures how well a language model predicts a text sample.\n\n    This metric is calculated as the average number of bits per word a model needs to represent the sample.\n\n    Args:\n        preds:\n            Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,\n            vocab_size], which is the output of a language model. Scores will be normalized internally using softmax.\n        target:\n            Ground truth values with a shape [batch_size, seq_len].\n        ignore_index:\n            Integer specifying a target class to ignore. If given, this class index does not contribute\n            to the returned score.\n\n    Returns:\n        Perplexity value\n\n    Examples:\n        >>> import torch\n        >>> gen = torch.manual_seed(42)\n        >>> preds = torch.rand(2, 8, 5, generator=gen)\n        >>> target = torch.randint(5, (2, 8), generator=gen)\n        >>> target[0, 6:] = -100\n        >>> perplexity(preds, target, ignore_index=-100)\n        tensor(5.8540)\n\n    \"\"\"\n    total, count = _perplexity_update(preds, target, ignore_index)\n    return _perplexity_compute(total, count)\n", "path": "src/torchmetrics/functional/text/perplexity.py"}], "after_files": [{"content": "# Copyright The Lightning team.\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\nfrom typing import Optional, Tuple\n\nimport torch\nfrom torch import Tensor\n\n\ndef _check_shape_and_type_consistency(preds: Tensor, target: Tensor) -> None:\n    \"\"\"Check shape and type consistency of input vectors.\n\n    Args:\n        preds:\n            Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,\n            vocab_size]. Scores will be normalized internally using softmax.\n        target:\n            Ground truth values with a shape [batch_size, seq_len].\n\n    Raises:\n        ValueError:\n            If ``preds`` tensor has no 3 dimensions.\n        ValueError:\n            If ``target`` tensor has no 2 dimensions.\n        ValueError:\n            If the first two dimensions of ``preds`` and ``target`` do not equal.\n        TypeError:\n            If ``preds`` dtype is not one of ``(torch.float16, torch.float32, torch.float64)``\n        TypeError:\n            If ``target`` is not of a type LongTensor (torch.int64)\n\n    \"\"\"\n    if len(preds.shape) != 3:\n        raise ValueError(\n            \"Input tensor `preds` is expected to have 3 dimensions, [batch_size, seq_len, vocab_size],\"\n            f\" but got {len(preds.shape)}.\"\n        )\n    if len(target.shape) != 2:\n        raise ValueError(\n            \"Input tensor `target` is expected to have 2 dimensions, [batch_size, seq_len],\"\n            f\" but got {len(target.shape)}.\"\n        )\n    if preds.shape[:2] != target.shape:\n        raise ValueError(\n            \"Input tensors `preds` and `target` are expected to have equaling first two dimensions,\"\n            f\" [batch_size, seq_len], but got {preds.shape[:2]} and {target.shape}.\"\n        )\n    if not preds.is_floating_point():\n        raise TypeError(f\"Input tensor `preds` is expected to be of floating point type but got {preds.dtype}.\")\n    if target.dtype != torch.int64:\n        raise TypeError(f\"Input tensor `target` is expected to be of a type {torch.int64} but got {target.dtype}.\")\n\n\ndef _perplexity_update(preds: Tensor, target: Tensor, ignore_index: Optional[int] = None) -> Tuple[Tensor, Tensor]:\n    \"\"\"Compute intermediate statistics for Perplexity.\n\n    Args:\n        preds:\n            Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,\n            vocab_size]. Scores will be normalized internally using softmax.\n        target:\n            Ground truth values with a shape [batch_size, seq_len].\n        ignore_index:\n            Integer specifying a target class to ignore. If given, this class index does not contribute\n            to the returned score.\n\n    Returns:\n        Log probabilities, summed over all samples\n        Number of samples\n\n    \"\"\"\n    _check_shape_and_type_consistency(preds, target)\n\n    probs = torch.nn.functional.softmax(preds.reshape(-1, preds.shape[-1]), dim=1)\n    target = target.reshape(-1)\n\n    if ignore_index is not None:\n        mask = target.ne(ignore_index)\n        target = target.where(target != ignore_index, torch.tensor(0, device=target.device))\n    else:\n        mask = torch.ones_like(target, dtype=torch.bool)\n\n    probs = probs[:, target].diagonal()[mask]\n    total_log_probs = -probs.log().sum()\n    count = mask.sum()\n\n    return total_log_probs, count\n\n\ndef _perplexity_compute(total: Tensor, count: Tensor) -> Tensor:\n    \"\"\"Compute the Perplexity.\n\n    Args:\n        total: Log probabilities, summed over all samples\n        count: Number of samples\n    Returns:\n        Perplexity\n\n    \"\"\"\n    return torch.exp(total / count)\n\n\ndef perplexity(preds: Tensor, target: Tensor, ignore_index: Optional[int] = None) -> Tensor:\n    \"\"\"Perplexity measures how well a language model predicts a text sample.\n\n    This metric is calculated as the average number of bits per word a model needs to represent the sample.\n\n    Args:\n        preds:\n            Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,\n            vocab_size], which is the output of a language model. Scores will be normalized internally using softmax.\n        target:\n            Ground truth values with a shape [batch_size, seq_len].\n        ignore_index:\n            Integer specifying a target class to ignore. If given, this class index does not contribute\n            to the returned score.\n\n    Returns:\n        Perplexity value\n\n    Examples:\n        >>> import torch\n        >>> gen = torch.manual_seed(42)\n        >>> preds = torch.rand(2, 8, 5, generator=gen)\n        >>> target = torch.randint(5, (2, 8), generator=gen)\n        >>> target[0, 6:] = -100\n        >>> perplexity(preds, target, ignore_index=-100)\n        tensor(5.8540)\n\n    \"\"\"\n    total, count = _perplexity_update(preds, target, ignore_index)\n    return _perplexity_compute(total, count)\n", "path": "src/torchmetrics/functional/text/perplexity.py"}]}