14637325149

Committed 24 Apr 2025 08:43AM UTC coverage: 84.932% (-3.2%) from 88.102%

Build # 14637325149

Build Type

Pull #295

github

Committed by

web-flow

Commit Message

Merge c93dc573f into a1f52e740

Pull Request Pull Request #295: Proposal to merge new integration of backends in PyManopt

Run Details

1915 of 2267 new or added lines in 31 files covered. (84.47%)

36 existing lines in 11 files now uncovered.

3534 of 4161 relevant lines covered (84.93%)

3.39 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

94.23

/src/pymanopt/function.py

__all__ = ["Function", "autograd", "jax", "numpy", "pytorch", "tensorflow"]

import inspect
from importlib import import_module
from typing import Any, Callable, Optional, Protocol

from pymanopt.backends import Backend
from pymanopt.manifolds.manifold import Manifold


class Function:
    def __init__(
        self, *, function: Callable, manifold: Manifold, backend: Backend
    ):
        if not callable(function):
            raise TypeError(f"Object {function} is not callable")

        self._original_function = function
        self._backend = backend
        self._function = backend.prepare_function(function)
        self._num_arguments = manifold.num_values

        self._gradient = None
        self._hessian = None

    def __str__(self):
        return f"Function <{self._backend}>"

    @property
    def backend(self):
        return self._backend

    def get_gradient_operator(self):
        if self._gradient is None:
            self._gradient = self._backend.generate_gradient_operator(
                self._original_function, self._num_arguments
            )
        return self._gradient

    def get_hessian_operator(self):
        if self._hessian is None:
            self._hessian = self._backend.generate_hessian_operator(
                self._original_function, self._num_arguments
            )
        return self._hessian

    def __call__(self, *args, **kwargs):
        return self._function(*args, **kwargs)


def _only_one_true(*args):
    return sum(args) == 1


class _ObjectiveFunctionDecorator(Protocol):
    def __call__(
        self, manifold: Manifold, dtype: Optional[Any] = None
    ) -> Callable[[Callable[..., Any]], Function]:
        ...


def decorator_factory(
    module: str, backend_class: str
) -> _ObjectiveFunctionDecorator:
    def decorator(
        manifold: Manifold, dtype: Optional[Any] = None
    ) -> Callable[[Callable[..., Any]], Function]:
        assert isinstance(manifold, Manifold)

        def inner(cost: Callable[..., Any]) -> Function:
            argspec = inspect.getfullargspec(cost)
            assert (
                _only_one_true(bool(argspec.args), bool(argspec.varargs))
                and not argspec.varkw
                and not argspec.kwonlyargs
            ), (
                "Decorated function must only accept positional arguments "
                "or a variable-length argument like *x"
            )
            backend_type = getattr(
                import_module(
                    f"pymanopt.backends.{module}",
                ),
                backend_class,
            )
            backend = (
                backend_type(dtype=dtype)
                if dtype is not None
                # by default use float64, which is fine for a function it only
                # uses autodiff methods (which do not depend on realness)
                else backend_type()
            )
            return Function(function=cost, manifold=manifold, backend=backend)

        return inner

    return decorator


numpy = decorator_factory("numpy_backend", "NumpyBackend")
jax = decorator_factory("jax_backend", "JaxBackend")
pytorch = decorator_factory("pytorch_backend", "PytorchBackend")
autograd = decorator_factory("autograd_backend", "AutogradBackend")
tensorflow = decorator_factory("tensorflow_backend", "TensorflowBackend")

1	__all__ = ["Function", "autograd", "jax", "numpy", "pytorch", "tensorflow"]	4✔
2
3	import inspect	4✔
4	from importlib import import_module	4✔
5	from typing import Any, Callable, Optional, Protocol	4✔
6
7	from pymanopt.backends import Backend	4✔
8	from pymanopt.manifolds.manifold import Manifold	4✔
9
10
11	class Function:	4✔
12	def __init__(	4✔
13	self, *, function: Callable, manifold: Manifold, backend: Backend
14	):
15	if not callable(function):	4✔
NEW 16	raise TypeError(f"Object {function} is not callable")	×
17
18	self._original_function = function	4✔
19	self._backend = backend	4✔
20	self._function = backend.prepare_function(function)	4✔
21	self._num_arguments = manifold.num_values	4✔
22
23	self._gradient = None	4✔
24	self._hessian = None	4✔
25
26	def __str__(self):	4✔
NEW 27	return f"Function <{self._backend}>"	×
28
29	@property	4✔
30	def backend(self):	4✔
31	return self._backend	4✔
32
33	def get_gradient_operator(self):	4✔
34	if self._gradient is None:	4✔
35	self._gradient = self._backend.generate_gradient_operator(	4✔
36	self._original_function, self._num_arguments
37	)
38	return self._gradient	4✔
39
40	def get_hessian_operator(self):	4✔
41	if self._hessian is None:	4✔
42	self._hessian = self._backend.generate_hessian_operator(	4✔
43	self._original_function, self._num_arguments
44	)
45	return self._hessian	4✔
46
47	def __call__(self, args, *kwargs):	4✔
48	return self._function(args, *kwargs)	4✔
49
50
51	def _only_one_true(*args):	4✔
52	return sum(args) == 1	4✔
53
54
55	class _ObjectiveFunctionDecorator(Protocol):	4✔
56	def __call__(	4✔
57	self, manifold: Manifold, dtype: Optional[Any] = None
58	) -> Callable[[Callable[..., Any]], Function]:
NEW 59	...	×
60
61
62	def decorator_factory(	4✔
63	module: str, backend_class: str
64	) -> _ObjectiveFunctionDecorator:
65	def decorator(	4✔
66	manifold: Manifold, dtype: Optional[Any] = None
67	) -> Callable[[Callable[..., Any]], Function]:
68	assert isinstance(manifold, Manifold)	4✔
69
70	def inner(cost: Callable[..., Any]) -> Function:	4✔
71	argspec = inspect.getfullargspec(cost)	4✔
72	assert (	4✔
73	_only_one_true(bool(argspec.args), bool(argspec.varargs))
74	and not argspec.varkw
75	and not argspec.kwonlyargs
76	), (
77	"Decorated function must only accept positional arguments "
78	"or a variable-length argument like *x"
79	)
80	backend_type = getattr(	4✔
81	import_module(
82	f"pymanopt.backends.{module}",
83	),
84	backend_class,
85	)
86	backend = (	4✔
87	backend_type(dtype=dtype)
88	if dtype is not None
89	# by default use float64, which is fine for a function it only
90	# uses autodiff methods (which do not depend on realness)
91	else backend_type()
92	)
93	return Function(function=cost, manifold=manifold, backend=backend)	4✔
94
95	return inner	4✔
96
97	return decorator	4✔
98
99
100	numpy = decorator_factory("numpy_backend", "NumpyBackend")	4✔
101	jax = decorator_factory("jax_backend", "JaxBackend")	4✔
102	pytorch = decorator_factory("pytorch_backend", "PytorchBackend")	4✔
103	autograd = decorator_factory("autograd_backend", "AutogradBackend")	4✔
104	tensorflow = decorator_factory("tensorflow_backend", "TensorflowBackend")	4✔

pymanopt / pymanopt / 14637325149

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous