internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: Lightning-AI/torchmetrics

Issue:
High memory usage of Perplexity metric
## 🐛 Bug

I ran out of memory (GPU) when computing the perplexity metric and would like to propose a small optimization to decrease its memory utilization.

### To Reproduce

For instance, when running the following code PyTorch tries to allocate 1024 GB of GPU memory on my system.
```py
from torchmetrics.text import Perplexity
import torch

gen = torch.manual_seed(42)
preds = torch.rand(512, 1024, 12, generator=gen).cuda()
target = torch.randint(12, (512, 1024), generator=gen).cuda()

perp = Perplexity().cuda()
print(perp(preds, target))
```

### Memory Inefficiency

I think the inefficiency is in this line:

https://github.com/Lightning-AI/torchmetrics/blob/a68455afb9041d1d32c1d6546897fee416abdc41/src/torchmetrics/functional/text/perplexity.py#L94

`probs[:, target]` results in a large temporary tensor with `(512*1024)^2` elements. Afterwards only the diagonal values are used. 

### Potential Solution

In contrast

```
probs = probs[torch.arange(target.numel()), target][mask]
```

would only require memory of the size of target. 

Would you consider accepting a pull request with this optimization? Or was the previous implementation chosen for another reason?

### Environment

- TorchMetrics v1.2.1 (installed with pip) and Master branch.
- Python 3.10.12
- Pytorch 2.2.0
- CUDA 12.1


Discussion:
Hi! thanks for your contribution!, great first issue!

Resolving patch:
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
@@ -91,7 +91,7 @@ def _perplexity_update(preds: Tensor, target: Tensor, ignore_index: Optional[int
     else:
         mask = torch.ones_like(target, dtype=torch.bool)
 
-    probs = probs[:, target].diagonal()[mask]
+    probs = probs[torch.arange(target.numel()), target][mask]
     total_log_probs = -probs.log().sum()
     count = mask.sum()
 


Test patch:
Repository: Lightning-AI/torchmetrics

Issue:
Can't access metrics in a MetricCollection via keys returned in MetricCollection.keys
## 🐛 Bug

print(list(self.val_metrics_macro.keys())) # prints "['val/MulticlassAccuracy/macro', 'val/MulticlassF1Score/macro', 'val/MulticlassPrecision/macro', 'val/MulticlassRecall/macro']"
print(
    f'val/accuracy/macro: {self.val_metrics_macro["val/MulticlassAccuracy/macro"].compute():.4}'
) # throws error: KeyError: 'val/MulticlassAccuracy/macro'

### Expected behavior

If a key shows up in `.keys` it should be a valid key in `__getitem__`.

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source): 1.3.1
- Python & PyTorch Version (e.g., 1.0): torch 2.2.0, python 3.10.12
- Any other relevant information such as OS (e.g., Linux): MacOS


Discussion:

Resolving patch:
diff --git a/src/torchmetrics/collections.py b/src/torchmetrics/collections.py
--- a/src/torchmetrics/collections.py
+++ b/src/torchmetrics/collections.py
@@ -547,6 +547,10 @@ def __getitem__(self, key: str, copy_state: bool = True) -> Metric:
 
         """
         self._compute_groups_create_state_ref(copy_state)
+        if self.prefix:
+            key = key.removeprefix(self.prefix)
+        if self.postfix:
+            key = key.removesuffix(self.postfix)
         return self._modules[key]
 
     @staticmethod


Test patch:
diff --git a/tests/unittests/bases/test_collections.py b/tests/unittests/bases/test_collections.py
--- a/tests/unittests/bases/test_collections.py
+++ b/tests/unittests/bases/test_collections.py
@@ -33,6 +33,7 @@
     MultilabelAveragePrecision,
 )
 from torchmetrics.utilities.checks import _allclose_recursive
+from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_2_0
 
 from unittests._helpers import seed_all
 from unittests._helpers.testers import DummyMetricDiff, DummyMetricMultiOutputDict, DummyMetricSum
@@ -150,6 +151,7 @@ def test_metric_collection_args_kwargs(tmpdir):
     assert metric_collection["DummyMetricDiff"].x == -20
 
 
+@pytest.mark.skipif(not _TORCH_GREATER_EQUAL_2_0, reason="Test requires torch 2.0 or higher")
 @pytest.mark.parametrize(
     ("prefix", "postfix"),
     [
@@ -204,6 +206,10 @@ def test_metric_collection_prefix_postfix_args(prefix, postfix):
     for name in names:
         assert f"new_prefix_{name}_new_postfix" in out, "postfix argument not working as intended with clone method"
 
+    keys = list(new_metric_collection.keys())
+    for k in keys:
+        assert new_metric_collection[k]  # check that the keys are valid even with prefix and postfix
+
 
 def test_metric_collection_repr():
     """Test MetricCollection."""
Repository: Lightning-AI/torchmetrics

Issue:
The default value of `compute_with_cache` should be `True` .
## 📚 Documentation

I found an inconsistency around `compute_with_cache` in `Metic`.
The docstring says the default value is `False` but the actual one is `True`.

https://github.com/Lightning-AI/torchmetrics/blob/ca7c8ebb7ee5ab63833cb2be4097e29b58b51b4d/src/torchmetrics/metric.py#L78
https://github.com/Lightning-AI/torchmetrics/blob/ca7c8ebb7ee5ab63833cb2be4097e29b58b51b4d/src/torchmetrics/metric.py#L140


Discussion:

Resolving patch:
diff --git a/src/torchmetrics/metric.py b/src/torchmetrics/metric.py
--- a/src/torchmetrics/metric.py
+++ b/src/torchmetrics/metric.py
@@ -75,7 +75,7 @@ class Metric(Module, ABC):
             - distributed_available_fn: Function that checks if the distributed backend is available. Defaults to a
               check of ``torch.distributed.is_available()`` and ``torch.distributed.is_initialized()``.
             - sync_on_compute: If metric state should synchronize when ``compute`` is called. Default is ``True``
-            - compute_with_cache: If results from ``compute`` should be cached. Default is ``False``
+            - compute_with_cache: If results from ``compute`` should be cached. Default is ``True``
 
     """
 


Test patch:
Repository: Lightning-AI/torchmetrics

Issue:
Documentation of ERGAS
## 📚 Documentation

It seems to me that the documentation of ERGAS are not consistent with the code. In the documentation the ratio h/l is multiplied with the rest, in the code it is a divisor. I think the confusion originates from the paper that is marked as the source. In the paper the ratio h/l < 1. The way the documentation states it, h/l should be > 1. Page 161 seems to give some insights on that.
The bug regarding these incosistencies was fixed in april, but now the documentation is not so clear and not consistent with the literature. 

Thanks!



Discussion:
Hi! thanks for your contribution!, great first issue!

Resolving patch:
diff --git a/src/torchmetrics/image/ergas.py b/src/torchmetrics/image/ergas.py
--- a/src/torchmetrics/image/ergas.py
+++ b/src/torchmetrics/image/ergas.py
@@ -35,12 +35,11 @@ class ErrorRelativeGlobalDimensionlessSynthesis(Metric):
     band of the result image. It is defined as:
 
     .. math::
-        ERGAS = 100 \cdot \frac{h}{l} \cdot \sqrt{\frac{1}{N} \sum_{k=1}^{N} \frac{RMSE(B_k)^2}{\mu_k^2}}
+        ERGAS = \frac{100}{r} \cdot \sqrt{\frac{1}{N} \sum_{k=1}^{N} \frac{RMSE(B_k)^2}{\mu_k^2}}
 
-    where :math:`h` and :math:`l` denote the spatial resolution (pixel size) of the high and low resolution images,
-    often shorted to the ratio between them :math:`r=h/l`. :math:`N` is the number of spectral bands, :math:`RMSE(B_k)`
-    is the root mean square error of the k-th band between low and high resolution images, and :math:`\\mu_k` is the
-    mean value of the k-th band of the reference image.
+    where :math:`r=h/l` denote the ratio in spatial resolution (pixel size) between the high and low resolution images.
+    :math:`N` is the number of spectral bands, :math:`RMSE(B_k)` is the root mean square error of the k-th band between
+    low and high resolution images, and :math:`\\mu_k` is the mean value of the k-th band of the reference image.
 
     As input to ``forward`` and ``update`` the metric accepts the following input
 


Test patch:
Repository: Lightning-AI/torchmetrics

Issue:
BootStrapper.reset() does not reset properly
## 🐛 Bug

Calling `BootStrapper.reset()` does not reset `self.metrics` properly.

### To Reproduce

```py
import torch
from torchmetrics.wrappers import BootStrapper
from torchmetrics.classification import MulticlassAccuracy

metric = BootStrapper(MulticlassAccuracy(num_classes=10))

for i in range(10):
    output = torch.randn((2000, 10))
    target = torch.randint(10, (2000,))
    # output = 0.5 * (target + output)
    metric.update(output, target)

print(metric.compute())
# {'mean': tensor(0.0990), 'std': tensor(0.0029)} <-- ok
print(metric.metrics[0].update_count)
# 10                                              <-- ok

metric.reset()

print(metric.compute())
# {'mean': tensor(0.0990), 'std': tensor(0.0029)} <-- ERROR, should be undefined after reset
print(metric.metrics[0].update_count)
# 10                                              <-- ERROR, should be 0 after reset

```

### Environment

- TorchMetrics version 1.4.0.post0
- Python version 3.11.9
- torch version 2.2.1+cu118



Discussion:

Resolving patch:
diff --git a/src/torchmetrics/wrappers/bootstrapping.py b/src/torchmetrics/wrappers/bootstrapping.py
--- a/src/torchmetrics/wrappers/bootstrapping.py
+++ b/src/torchmetrics/wrappers/bootstrapping.py
@@ -168,6 +168,12 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
         """Use the original forward method of the base metric class."""
         return super(WrapperMetric, self).forward(*args, **kwargs)
 
+    def reset(self) -> None:
+        """Reset the state of the base metric."""
+        for m in self.metrics:
+            m.reset()
+        super().reset()
+
     def plot(
         self, val: Optional[Union[Tensor, Sequence[Tensor]]] = None, ax: Optional[_AX_TYPE] = None
     ) -> _PLOT_OUT_TYPE:


Test patch:
diff --git a/tests/unittests/wrappers/test_bootstrapping.py b/tests/unittests/wrappers/test_bootstrapping.py
--- a/tests/unittests/wrappers/test_bootstrapping.py
+++ b/tests/unittests/wrappers/test_bootstrapping.py
@@ -124,6 +124,17 @@ def test_bootstrap(device, sampling_strategy, metric, ref_metric):
     assert np.allclose(output["std"].cpu(), np.std(sk_scores, ddof=1))
     assert np.allclose(output["raw"].cpu(), sk_scores)
 
+    # check that resetting works
+    bootstrapper.reset()
+
+    assert bootstrapper.update_count == 0
+    assert all(m.update_count == 0 for m in bootstrapper.metrics)
+    output = bootstrapper.compute()
+    if not isinstance(metric, MeanSquaredError):
+        assert output["mean"] == 0
+        assert output["std"] == 0
+        assert (output["raw"] == torch.zeros(10, device=device)).all()
+
 
 @pytest.mark.parametrize("sampling_strategy", ["poisson", "multinomial"])
 def test_low_sample_amount(sampling_strategy):
Repository: google/TensorNetwork

Issue:
Python version requirement
Our current setup file is still happy with python>=3.5, whereas we stopped testing for it and it seems that Jax also discontinued support. Should we not be stricter in the setup especially if we move to Jax as default?


Discussion:
Yes. Please make a PR as we'll want that for the next release.

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -36,7 +36,7 @@
     url='http://github.com/google/TensorNetwork',
     author='The TensorNetwork Developers',
     author_email='chaseriley@google.com',
-    python_requires=('>=3.5.0'),
+    python_requires=('>=3.6.0'),
     install_requires=requirements,
     license='Apache 2.0',
     description=description,


Test patch:
Repository: google/TensorNetwork

Issue:
Pytorch 1.4.0 breaks our build
Needs to be investigated.


Discussion:

Resolving patch:
diff --git a/tensornetwork/backends/pytorch/pytorch_backend.py b/tensornetwork/backends/pytorch/pytorch_backend.py
--- a/tensornetwork/backends/pytorch/pytorch_backend.py
+++ b/tensornetwork/backends/pytorch/pytorch_backend.py
@@ -249,7 +249,7 @@ def eigsh_lanczos(self,
             self.torch.tensor(diag_elements)) + self.torch.diag(
                 self.torch.tensor(norms_vector_n[1:]), 1) + self.torch.diag(
                     self.torch.tensor(norms_vector_n[1:]), -1)
-        eigvals, u = A_tridiag.symeig()
+        eigvals, u = A_tridiag.symeig(eigenvectors=True)
         if not first:
           if self.torch.norm(eigvals[0:numeig] - eigvalsold[0:numeig]) < tol:
             break
@@ -266,7 +266,7 @@ def eigsh_lanczos(self,
         self.torch.tensor(diag_elements)) + self.torch.diag(
             self.torch.tensor(norms_vector_n[1:]), 1) + self.torch.diag(
                 self.torch.tensor(norms_vector_n[1:]), -1)
-    eigvals, u = A_tridiag.symeig()
+    eigvals, u = A_tridiag.symeig(eigenvectors=True)
     eigenvectors = []
     for n2 in range(min(numeig, len(eigvals))):
       state = self.zeros(initial_state.shape, initial_state.dtype)


Test patch:
diff --git a/tensornetwork/backends/pytorch/pytorch_backend_test.py b/tensornetwork/backends/pytorch/pytorch_backend_test.py
--- a/tensornetwork/backends/pytorch/pytorch_backend_test.py
+++ b/tensornetwork/backends/pytorch/pytorch_backend_test.py
@@ -247,7 +247,7 @@ def mv(x):
     return H.mv(x)
 
   eta1, U1 = backend.eigsh_lanczos(mv, init)
-  eta2, U2 = H.symeig()
+  eta2, U2 = H.symeig(eigenvectors=True)
   v2 = U2[:, 0]
   v2 = v2 / sum(v2)
   v1 = np.reshape(U1[0], (D))
@@ -275,7 +275,7 @@ def __call__(self, x):
 
   mv = LinearOperator(shape=((D,), (D,)), dtype=dtype)
   eta1, U1 = backend.eigsh_lanczos(mv, init)
-  eta2, U2 = H.symeig()
+  eta2, U2 = H.symeig(eigenvectors=True)
   v2 = U2[:, 0]
   v2 = v2 / sum(v2)
   v1 = np.reshape(U1[0], (D))
@@ -303,7 +303,7 @@ def __call__(self, x):
   mv = LinearOperator(shape=((D,), (D,)), dtype=dtype)
   eta1, U1 = backend.eigsh_lanczos(mv, reorthogonalize=True, ndiag=1,
                                    tol=10**(-12), delta=10**(-12))
-  eta2, U2 = H.symeig()
+  eta2, U2 = H.symeig(eigenvectors=True)
   v2 = U2[:, 0]
   v2 = v2 / sum(v2)
   v1 = np.reshape(U1[0], (D))
Repository: google/TensorNetwork

Issue:
tn.set_default_backend should raise exception
`tn.set_default_backend(backend_name)` should raise if `backend_name` is not a valid backend.


Discussion:
Yes it should!

Resolving patch:
diff --git a/tensornetwork/backend_contextmanager.py b/tensornetwork/backend_contextmanager.py
--- a/tensornetwork/backend_contextmanager.py
+++ b/tensornetwork/backend_contextmanager.py
@@ -1,5 +1,6 @@
 from typing import Text, Union
 from tensornetwork.backends.base_backend import BaseBackend
+from tensornetwork.backends import backend_factory
 
 class DefaultBackend():
   """Context manager for setting up backend for nodes"""
@@ -38,4 +39,6 @@ def set_default_backend(backend: Union[Text, BaseBackend]) -> None:
   if not isinstance(backend, (Text, BaseBackend)):
     raise ValueError("Item passed to set_default_backend "
                      "must be Text or BaseBackend")
+  if isinstance(backend, Text) and backend not in backend_factory._BACKENDS:
+    raise ValueError(f"Backend '{backend}' was not found.")
   _default_backend_stack.default_backend = backend


Test patch:
diff --git a/tensornetwork/tests/tensornetwork_test.py b/tensornetwork/tests/tensornetwork_test.py
--- a/tensornetwork/tests/tensornetwork_test.py
+++ b/tensornetwork/tests/tensornetwork_test.py
@@ -582,6 +582,10 @@ def test_set_default(backend):
   a = tn.Node(np.eye(2))
   assert a.backend.name == backend
 
+def test_bad_backend_name():
+  with pytest.raises(ValueError, match="Backend 'BAD_NAME' was not found."):
+    tn.set_default_backend("BAD_NAME")
+
 def test_copy_tensor(backend):
   a = tn.Node(np.array([1, 2, 3], dtype=np.float64), backend=backend)
   b = tn.Node(np.array([10, 20, 30], dtype=np.float64), backend=backend)
Repository: google/TensorNetwork

Issue:
Implement get_all_dangling_ordered
Seems like new users get confused by `get_all_dangling` returning a `set` instead of a `list`. We can mitigate this by creating a `get_all_dangling_ordered` function.

See #481 as an example.


Discussion:

Resolving patch:
diff --git a/tensornetwork/network_components.py b/tensornetwork/network_components.py
--- a/tensornetwork/network_components.py
+++ b/tensornetwork/network_components.py
@@ -331,9 +331,9 @@ def get_all_nondangling(self) -> Set["Edge"]:
     """Return the set of nondangling edges connected to this node."""
     return {edge for edge in self.edges if not edge.is_dangling()}
 
-  def get_all_dangling(self) -> Set["Edge"]:
+  def get_all_dangling(self) -> List["Edge"]:
     """Return the set of dangling edges connected to this node."""
-    return {edge for edge in self.edges if edge.is_dangling()}
+    return [edge for edge in self.edges if edge.is_dangling()]
 
   def set_name(self, name) -> None:
     if not isinstance(name, str):
@@ -1311,11 +1311,11 @@ def get_all_nondangling(nodes: Iterable[BaseNode]) -> Set[Edge]:
   return edges
 
 
-def get_all_dangling(nodes: Iterable[BaseNode]) -> Set[Edge]:
+def get_all_dangling(nodes: Iterable[BaseNode]) -> List[Edge]:
   """Return the set of all dangling edges."""
-  edges = set()
+  edges = []
   for node in nodes:
-    edges |= node.get_all_dangling()
+    edges += node.get_all_dangling()
   return edges
 
 


Test patch:
Repository: google/TensorNetwork

Issue:
Some doc links to github yield error
The "Edit on GitHub" links in the top right of some doc pages yield 404 errors when followed, for example [tn.Node](https://tensornetwork.readthedocs.io/en/latest/stubs/tensornetwork.Node.html) and [tn.contractors.optimal](https://tensornetwork.readthedocs.io/en/latest/stubs/tensornetwork.contractors.optimal.html#tensornetwork.contractors.optimal).

The links at the top of list pages work, for example [common functions](https://tensornetwork.readthedocs.io/en/latest/network.html) and [contractors](https://tensornetwork.readthedocs.io/en/latest/contractors.html).

Possibly those are meant to be links to edit the doc pages themselves, rather than the source code?


Discussion:
Huh, weird. The Edit on Github link is auto generated, so I'd have to investigate how that is setup. 

Resolving patch:
diff --git a/docs/conf.py b/docs/conf.py
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -43,6 +43,7 @@
 exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
 
 autosummary_generate = True
+autosummary_generate_overwrite = False
 napolean_use_rtype = False
 # -- Options for HTML output -------------------------------------------------
 


Test patch:
Repository: google/TensorNetwork

Issue:
add dynamic programming contractor of opt_einsum
We should update this asap


Discussion:
Yes we should.

Should we just change the `contractors.optimal` to use `optimize='dp'`? Or should we add a `tn.contractors.dp` method? I'm ok with either, with slight preference to the former.
I think the first choice is better, since it seems that 'dp' is destined to become the new standard in this case anyway at some point.

Resolving patch:
diff --git a/tensornetwork/contractors/opt_einsum_paths/path_contractors.py b/tensornetwork/contractors/opt_einsum_paths/path_contractors.py
--- a/tensornetwork/contractors/opt_einsum_paths/path_contractors.py
+++ b/tensornetwork/contractors/opt_einsum_paths/path_contractors.py
@@ -117,7 +117,8 @@ def optimal(nodes: Iterable[AbstractNode],
   Returns:
     The final node after full contraction.
   """
-  alg = functools.partial(opt_einsum.paths.optimal, memory_limit=memory_limit)
+  alg = functools.partial(
+      opt_einsum.paths.dynamic_programming, memory_limit=memory_limit)
   return base(nodes, alg, output_edge_order, ignore_edge_order)
 
 


Test patch:
Repository: google/TensorNetwork

Issue:
bump tests to python 3.8



Discussion:

Resolving patch:
diff --git a/tensornetwork/block_sparse/blocksparse_utils.py b/tensornetwork/block_sparse/blocksparse_utils.py
--- a/tensornetwork/block_sparse/blocksparse_utils.py
+++ b/tensornetwork/block_sparse/blocksparse_utils.py
@@ -597,7 +597,7 @@ def _find_transposed_diagonal_sparse_blocks(
     return cacher.cache[hash_val]
   return block_maps, obj, block_dims
 
-def _to_string(charges: List[BaseCharge], flows: np.ndarray,
+def _to_string(charges: List[BaseCharge], flows: Union[np.ndarray, List],
                tr_partition: int, order: List[int]) -> str:
   """
   map the input arguments of _find_transposed_diagonal_sparse_blocks 
diff --git a/tensornetwork/network_operations.py b/tensornetwork/network_operations.py
--- a/tensornetwork/network_operations.py
+++ b/tensornetwork/network_operations.py
@@ -604,7 +604,7 @@ def reachable(
   if isinstance(inputs, AbstractNode):
     inputs = {inputs}
   elif isinstance(inputs, Edge):
-    inputs = {inputs.node1}
+    inputs = {inputs.node1} # pytype: disable=attribute-error
   elif isinstance(inputs, list) and all(isinstance(x, Edge) for x in inputs):
     inputs = {x.node1 for x in inputs}
   return _reachable(set(inputs))
@@ -793,7 +793,7 @@ def switch_backend(nodes: Iterable[AbstractNode], new_backend: Text) -> None:
     node.backend = backend
 
 
-def get_neighbors(node: AbstractNode) -> List[Node]:
+def get_neighbors(node: AbstractNode) -> List[AbstractNode]:
   """Get all of the neighbors that are directly connected to the given node.
 
   Note: `node` will never be in the returned list, even if `node` has a


Test patch:
Repository: google/TensorNetwork

Issue:
Latest JAX breaks build
Looks like it's mainly problems related to using `jnp` for shape calculations.


Discussion:
I'll have a look. Cleaning up jax backend right now 

Resolving patch:
diff --git a/tensornetwork/backends/jax/jitted_functions.py b/tensornetwork/backends/jax/jitted_functions.py
--- a/tensornetwork/backends/jax/jitted_functions.py
+++ b/tensornetwork/backends/jax/jitted_functions.py
@@ -108,8 +108,8 @@ def check_thresh(check_vals):
 
       return jax.lax.cond(iteration <= maxiteration, [norm, threshold],
                           check_thresh, False, lambda x: x)
-
-    numel = jax.numpy.prod(init.shape)
+    #TODO (mganahl): check if this runs on TPU (dtype issue)
+    numel = np.prod(init.shape).astype(np.int32)
     krylov_vecs = jax.numpy.zeros((ncv + 1, numel), dtype=init.dtype)
     norms = jax.numpy.zeros(ncv, dtype=init.dtype)
     diag_elems = jax.numpy.zeros(ncv, dtype=init.dtype)


Test patch: