facebookresearch / facebookresearch/theseus
Vectorization can give incorrect results if similar costs depend on values not declared as optim/aux vars
- Dominant language
- Python
- Stars
- 2.1k
- Forks
- 149
- PR merge metrics
- No merged PRs in 30d
Description
## 🐛 Bug
## Steps to Reproduce
Here is a simple repro
```python
import theseus as th
import torch
from theseus.core.vectorizer import Vectorize
class BadCost(th.CostFunction):
def __init__(self, x, flag=False, name=None):
super().__init__(th.ScaleCostWeight(1.0), name=name)
self.x = x
self.flag = flag
self.register_optim_vars(["x"])
def error(self) -> torch.Tensor:
return (
torch.ones_like(self.x.tensor)
if self.flag
else torch.zeros_like(self.x.tensor)
)
def jacobians(self):
raise NotImplementedError
def dim(self) -> int:
return self.x.dof()
def _copy_impl(self, new_name=None):
return BadCost(self.x.copy(), flag=self.flag, name=new_name) # type: ignore
o = th.Objective()
z = th.Vector(1, name="z")
o.add(BadCost(z, flag=True))
o.add(BadCost(z, flag=False))
o.update() # to resolve batch size for vectorization
Vectorize(o)
print(o.error())
```
## Expected behavior
The above prints `tensor([[1., 1.]])` but it should print `tensor([[1., 0.]])`.
Contributor guide
Assessment
This issue has not been assessed yet.