rfeinman / pytorch-minimize

Newton and Quasi-Newton optimization with PyTorch
https://pytorch-minimize.readthedocs.io
MIT License
292 stars 34 forks source link

make `minimize` jit-able? #4

Open alibool opened 2 years ago

alibool commented 2 years ago

Thanks for this great package, it seems thatminimize is not jit-able? for example

from torchmin import minimize
import torch
from functools import partial

def bar(y, x):
    return (x + y).square().sum()

# raise error
@torch.jit.script
def foo(x: torch.Tensor) -> torch.Tensor:
    result = minimize(partial(bar, x=x), torch.tensor([0., 1.]), method='bfgs')
    return result.x

Will you consider making minimize jit-able?

rfeinman commented 2 years ago

Hi @alibool - thanks for your comment and suggestion.

This would require a significant refactor. I'd have to step back through each minimization routine, add type hints to the function signatures and make sure that all torch ops are TorchScript-supported. It's on my agenda but it won't be a high priority (unless I have some help).