Commit 70a698c2 authored by Philipp Arras's avatar Philipp Arras
Browse files

Cosmetics

parent 80fd5456
...@@ -15,11 +15,10 @@ ...@@ -15,11 +15,10 @@
# Author: Philipp Arras # Author: Philipp Arras
import numpy as np import numpy as np
from .operator import Operator
from .energy_operators import EnergyOperator, LikelihoodEnergyOperator
from .linear_operator import LinearOperator
from .endomorphic_operator import EndomorphicOperator
from .energy_operators import LikelihoodEnergyOperator
from .linear_operator import LinearOperator
from .operator import Operator
try: try:
import jax import jax
...@@ -60,8 +59,8 @@ class JaxOperator(Operator): ...@@ -60,8 +59,8 @@ class JaxOperator(Operator):
self._fwd = jax.jit(lambda x, y: jax.jvp(self._func, (x,), (y,))[1]) self._fwd = jax.jit(lambda x, y: jax.jvp(self._func, (x,), (y,))[1])
def apply(self, x): def apply(self, x):
from ..sugar import is_linearization, makeField
from ..multi_domain import MultiDomain from ..multi_domain import MultiDomain
from ..sugar import is_linearization, makeField
self._check_input(x) self._check_input(x)
if is_linearization(x): if is_linearization(x):
res, bwd = self._vjp(x.val.val) res, bwd = self._vjp(x.val.val)
...@@ -137,9 +136,9 @@ class JaxLikelihoodEnergyOperator(LikelihoodEnergyOperator): ...@@ -137,9 +136,9 @@ class JaxLikelihoodEnergyOperator(LikelihoodEnergyOperator):
return self._dt, self._trafo return self._dt, self._trafo
def apply(self, x): def apply(self, x):
from ..linearization import Linearization
from ..sugar import is_linearization, makeField from ..sugar import is_linearization, makeField
from .simple_linear_operators import VdotOperator from .simple_linear_operators import VdotOperator
from ..linearization import Linearization
self._check_input(x) self._check_input(x)
lin = is_linearization(x) lin = is_linearization(x)
val = x.val.val if lin else x.val val = x.val.val if lin else x.val
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment