Reinstate black-compatible isort (#450)

* Reinstate isort

* Run isort

* Fix ops.AssociativeOp issue

* Ignore cuda warning during tests

* Attempt to fix pattern registration

* Another attempt

* Fix merge conflict

Co-authored-by: eb8680 <[email protected]>

Makefile

@@ -12,12 +12,14 @@ docs: FORCE
 lint: FORCE
 	flake8
 	black --check .
+	isort --check .
 
 license: FORCE
 	python scripts/update_headers.py
 
 format: FORCE
 	black .
+	isort .
 
 test: lint FORCE
 ifeq (${FUNSOR_BACKEND}, torch)

examples/discrete_hmm.py

@@ -7,8 +7,8 @@
 import torch
 
 import funsor
-import funsor.torch.distributions as dist
 import funsor.ops as ops
+import funsor.torch.distributions as dist
 from funsor.interpreter import interpretation, reinterpret
 from funsor.optimizer import apply_optimizer
 from funsor.terms import lazy

examples/eeg_slds.py

@@ -17,13 +17,13 @@
 from urllib.request import urlopen
 
 import numpy as np
+import pyro
 import torch
 import torch.nn as nn
-import pyro
 
 import funsor
-import funsor.torch.distributions as dist
 import funsor.ops as ops
+import funsor.torch.distributions as dist
 from funsor.pyro.convert import (
     funsor_to_cat_and_mvn,
     funsor_to_mvn,

examples/kalman_filter.py

@@ -6,8 +6,8 @@
 import torch
 
 import funsor
-import funsor.torch.distributions as dist
 import funsor.ops as ops
+import funsor.torch.distributions as dist
 from funsor.interpreter import interpretation, reinterpret
 from funsor.optimizer import apply_optimizer
 from funsor.terms import lazy

examples/mixed_hmm/experiment.py

@@ -10,15 +10,15 @@
 import pyro
 import pyro.poutine as poutine
 import torch
+from model import Guide, Model
+from seal_data import prepare_fake, prepare_seal
 
 import funsor
 import funsor.ops as ops
 from funsor.interpreter import interpretation
 from funsor.optimizer import apply_optimizer
 from funsor.sum_product import MarkovProduct, naive_sequential_sum_product, sum_product
 from funsor.terms import lazy, to_funsor
-from model import Guide, Model
-from seal_data import prepare_fake, prepare_seal
 
 
 def aic_num_parameters(model, guide=None):

examples/mixed_hmm/model.py

@@ -7,8 +7,8 @@
 import torch
 from torch.distributions import constraints
 
-import funsor.torch.distributions as dist
 import funsor.ops as ops
+import funsor.torch.distributions as dist
 from funsor.domains import Bint, Reals
 from funsor.tensor import Tensor
 from funsor.terms import Stack, Variable, to_funsor

examples/sensor.py

@@ -12,8 +12,8 @@
 from torch.optim import Adam
 
 import funsor
-import funsor.torch.distributions as f_dist
 import funsor.ops as ops
+import funsor.torch.distributions as f_dist
 from funsor.domains import Reals
 from funsor.pyro.convert import dist_to_funsor, funsor_to_mvn
 from funsor.tensor import Tensor, Variable
@@ -241,8 +241,8 @@ def main(args):
         import matplotlib
 
         matplotlib.use("Agg")
-        from matplotlib import pyplot
         import numpy as np
+        from matplotlib import pyplot
 
         seeds = set(seed for seed, _, _ in results)
         X = args.num_frames

examples/slds.py

@@ -6,8 +6,8 @@
 import torch
 
 import funsor
-import funsor.torch.distributions as dist
 import funsor.ops as ops
+import funsor.torch.distributions as dist
 
 
 def main(args):

examples/vae.py

@@ -12,8 +12,8 @@
 from torchvision import datasets, transforms
 
 import funsor
-import funsor.torch.distributions as dist
 import funsor.ops as ops
+import funsor.torch.distributions as dist
 from funsor.domains import Bint, Reals
 
 REPO_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

funsor/__init__.py

@@ -3,7 +3,7 @@
 
 from funsor.domains import Array, Bint, Domain, Real, Reals, bint, find_domain, reals
 from funsor.integrate import Integrate
-from funsor.interpreter import reinterpret, interpretation
+from funsor.interpreter import interpretation, reinterpret
 from funsor.sum_product import MarkovProduct
 from funsor.tensor import Tensor, function
 from funsor.terms import (
@@ -42,7 +42,6 @@
     testing,
 )
 
-
 __all__ = [
     "Array",
     "Bint",

funsor/affine.py

@@ -8,7 +8,7 @@
 
 from funsor.interpreter import gensym
 from funsor.tensor import Einsum, Tensor, get_default_prototype
-from funsor.terms import Binary, Funsor, Lambda, Reduce, Unary, Variable, Bint
+from funsor.terms import Binary, Bint, Funsor, Lambda, Reduce, Unary, Variable
 
 from . import ops
 

funsor/distribution.py

@@ -43,7 +43,6 @@
 )
 from funsor.util import broadcast_shape, get_backend, getargspec, lazy_property
 
-
 BACKEND_TO_DISTRIBUTIONS_BACKEND = {
     "torch": "funsor.torch.distributions",
     "jax": "funsor.jax.distributions",

funsor/einsum/numpy_map.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import operator
-
 from functools import reduce
 
 import funsor.ops as ops

funsor/integrate.py

@@ -7,7 +7,7 @@
 import funsor.ops as ops
 from funsor.cnf import Contraction, GaussianMixture
 from funsor.delta import Delta
-from funsor.gaussian import Gaussian, align_gaussian, _mv, _trace_mm, _vv
+from funsor.gaussian import Gaussian, _mv, _trace_mm, _vv, align_gaussian
 from funsor.tensor import Tensor
 from funsor.terms import (
     Funsor,

funsor/jax/__init__.py

@@ -4,20 +4,23 @@
 from jax.core import Tracer
 from jax.interpreters.xla import DeviceArray
 
-import funsor.jax.distributions  # noqa: F401
-import funsor.jax.ops  # noqa: F401
-import funsor.ops as ops
 from funsor.adjoint import adjoint_ops
 from funsor.interpreter import children, recursion_reinterpret
-from funsor.terms import Funsor, to_funsor
+from funsor.ops import AssociativeOp
 from funsor.tensor import Tensor, tensor_to_funsor
+from funsor.terms import Funsor, to_funsor
 from funsor.util import quote
 
+from . import distributions as _
+from . import ops as _
+
+del _  # flake8
+
 
 @adjoint_ops.register(
     Tensor,
-    ops.AssociativeOp,
-    ops.AssociativeOp,
+    AssociativeOp,
+    AssociativeOp,
     Funsor,
     (DeviceArray, Tracer),
     tuple,

funsor/jax/distributions.py

@@ -6,10 +6,11 @@
 
 import numpyro.distributions as dist
 
+import funsor.ops as ops
 from funsor.cnf import Contraction
 from funsor.distribution import (
-    Bernoulli,
     FUNSOR_DIST_NAMES,
+    Bernoulli,
     LogNormal,
     backenddist_to_funsor,
     eager_beta,
@@ -37,7 +38,6 @@
     transformeddist_to_funsor,
 )
 from funsor.domains import Real, Reals
-import funsor.ops as ops
 from funsor.tensor import Tensor
 from funsor.terms import Binary, Funsor, Reduce, Variable, eager, to_data, to_funsor
 from funsor.util import methodof

funsor/jax/ops.py

@@ -11,7 +11,7 @@
 from jax.scipy.linalg import cho_solve, solve_triangular
 from jax.scipy.special import expit, gammaln, logsumexp
 
-import funsor.ops as ops
+from .. import ops
 
 ################################################################################
 # Register Ops

funsor/ops/__init__.py

@@ -1,6 +1,7 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+from . import array, builtin, op
 from .array import *
 from .builtin import *
 from .op import *

funsor/ops/builtin.py

@@ -181,6 +181,9 @@ def sigmoid_log_abs_det_jacobian(x, y):
 PRODUCT_INVERSES[add] = sub
 
 __all__ = [
+    "AssociativeOp",
+    "GetitemOp",
+    "NullOp",
     "abs",
     "add",
     "and_",

funsor/pyro/distribution.py

@@ -4,8 +4,8 @@
 from collections import OrderedDict
 
 import torch
-from torch.distributions import constraints
 from pyro.distributions import TorchDistribution
+from torch.distributions import constraints
 
 from funsor.cnf import Contraction
 from funsor.delta import Delta

funsor/tensor.py

@@ -15,11 +15,12 @@
 from multipledispatch.variadic import Variadic
 
 import funsor
-import funsor.ops as ops
-from funsor.delta import Delta
-from funsor.domains import Array, ArrayType, Bint, Product, Real, Reals, find_domain
-from funsor.ops import GetitemOp, MatmulOp, Op, ReshapeOp
-from funsor.terms import (
+
+from . import ops
+from .delta import Delta
+from .domains import Array, ArrayType, Bint, Product, Real, Reals, find_domain
+from .ops import GetitemOp, MatmulOp, Op, ReshapeOp
+from .terms import (
     Binary,
     Funsor,
     FunsorMeta,
@@ -34,7 +35,7 @@
     to_data,
     to_funsor,
 )
-from funsor.util import get_backend, get_tracing_state, getargspec, is_nn_module, quote
+from .util import get_backend, get_tracing_state, getargspec, is_nn_module, quote
 
 
 def get_default_prototype():

funsor/testing.py

@@ -17,10 +17,10 @@
 import funsor.ops as ops
 from funsor.cnf import Contraction
 from funsor.delta import Delta
-from funsor.domains import Domain, Bint, Real
+from funsor.domains import Bint, Domain, Real
 from funsor.gaussian import Gaussian
-from funsor.terms import Funsor, Number
 from funsor.tensor import Tensor
+from funsor.terms import Funsor, Number
 from funsor.util import get_backend
 
 

funsor/torch/__init__.py

@@ -4,18 +4,21 @@
 import torch
 from multipledispatch import dispatch
 
-import funsor.torch.distributions  # noqa: F401
-import funsor.torch.ops  # noqa: F401
-import funsor.ops as ops
 from funsor.adjoint import adjoint_ops
 from funsor.interpreter import children, recursion_reinterpret
-from funsor.terms import Funsor, to_funsor
+from funsor.ops import AssociativeOp
 from funsor.tensor import Tensor, tensor_to_funsor
+from funsor.terms import Funsor, to_funsor
 from funsor.util import quote
 
+from . import distributions as _
+from . import ops as _
+
+del _  # flake8
+
 
 @adjoint_ops.register(
-    Tensor, ops.AssociativeOp, ops.AssociativeOp, Funsor, torch.Tensor, tuple, object
+    Tensor, AssociativeOp, AssociativeOp, Funsor, torch.Tensor, tuple, object
 )
 def adjoint_tensor(adj_redop, adj_binop, out_adj, data, inputs, dtype):
     return {}

funsor/util.py

@@ -178,12 +178,14 @@ def set_backend(backend):
         _FUNSOR_BACKEND = "torch"
 
         import torch  # noqa: F401
+
         import funsor.torch  # noqa: F401
     elif backend == "jax":
         _FUNSOR_BACKEND = "jax"
         _JAX_LOADED = True
 
         import jax  # noqa: F401
+
         import funsor.jax  # noqa: F401
     else:
         raise ValueError(

scripts/update_headers.py

@@ -1,8 +1,8 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
-import os
 import glob
+import os
 
 root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 blacklist = ["/build/", "/dist/"]

setup.cfg

@@ -9,6 +9,11 @@ per-file-ignores =
     funsor/jax/distributions.py:F821
     funsor/torch/distributions.py:F821
 
+[isort]
+profile = black
+known_first_party = funsor, test
+known_third_party = opt_einsum, pyro, pyroapi, torch, torchvision
+
 [tool:pytest]
 filterwarnings = error
     ignore:numpy.ufunc size changed:RuntimeWarning

setup.py

@@ -51,6 +51,7 @@
         "test": [
             "black",
             "flake8",
+            "isort>=5.0",
             "pandas",
             "pyro-api>=0.1.2",
             "pytest==4.3.1",
@@ -61,6 +62,7 @@
         "dev": [
             "black",
             "flake8",
+            "isort>=5.0",
             "pandas",
             "pytest==4.3.1",
             "pytest-xdist==1.27.0",

test/examples/test_bart.py

@@ -7,8 +7,8 @@
 import torch
 
 import funsor
-import funsor.torch.distributions as dist
 import funsor.ops as ops
+import funsor.torch.distributions as dist
 from funsor.cnf import Contraction
 from funsor.domains import Bint, Real, Reals
 from funsor.gaussian import Gaussian

test/examples/test_sensor_fusion.py

@@ -6,8 +6,8 @@
 import pytest
 import torch
 
-import funsor.torch.distributions as dist
 import funsor.ops as ops
+import funsor.torch.distributions as dist
 from funsor.cnf import Contraction
 from funsor.domains import Bint, Reals
 from funsor.gaussian import Gaussian