From a9799a38b810c8ac7a280521f4c09910541e26fd Mon Sep 17 00:00:00 2001 From: pranmod01 Date: Fri, 18 Oct 2024 09:04:37 -0700 Subject: [PATCH] isort fix --- jax | 1 - src/nemos/__init__.py | 2 +- src/nemos/_documentation_utils/__init__.py | 2 +- src/nemos/base_regressor.py | 5 ++--- src/nemos/convolve.py | 3 +-- src/nemos/glm.py | 5 ++--- src/nemos/observation_models.py | 3 +-- src/nemos/regularizer.py | 5 ++--- src/nemos/simulation.py | 5 ++--- src/nemos/solvers.py | 7 +++---- src/nemos/type_casting.py | 5 ++--- src/nemos/typing.py | 3 +-- src/nemos/utils.py | 5 ++--- src/nemos/validation.py | 3 +-- 14 files changed, 21 insertions(+), 33 deletions(-) delete mode 160000 jax diff --git a/jax b/jax deleted file mode 160000 index 9cf952a5..00000000 --- a/jax +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 9cf952a535518da59cdcecc9145dba287beddca2 diff --git a/src/nemos/__init__.py b/src/nemos/__init__.py index 97c5b3db..aedd05c0 100644 --- a/src/nemos/__init__.py +++ b/src/nemos/__init__.py @@ -14,5 +14,5 @@ styles, tree_utils, type_casting, - utils, + utils ) diff --git a/src/nemos/_documentation_utils/__init__.py b/src/nemos/_documentation_utils/__init__.py index 3cd63e0e..1c64a43a 100644 --- a/src/nemos/_documentation_utils/__init__.py +++ b/src/nemos/_documentation_utils/__init__.py @@ -19,5 +19,5 @@ plot_rates_and_smoothed_counts, plot_weighted_sum_basis, run_animation, - tuning_curve_plot, + tuning_curve_plot ) diff --git a/src/nemos/base_regressor.py b/src/nemos/base_regressor.py index ba2782ed..e4a425ce 100644 --- a/src/nemos/base_regressor.py +++ b/src/nemos/base_regressor.py @@ -9,11 +9,10 @@ from copy import deepcopy from typing import Any, Dict, NamedTuple, Optional, Tuple, Union -import jaxopt -from numpy.typing import ArrayLike, NDArray - import jax import jax.numpy as jnp +import jaxopt +from numpy.typing import ArrayLike, NDArray from . import solvers, utils, validation from ._regularizer_builder import AVAILABLE_REGULARIZERS, create_regularizer diff --git a/src/nemos/convolve.py b/src/nemos/convolve.py index e9a8d1fc..6a6294c3 100644 --- a/src/nemos/convolve.py +++ b/src/nemos/convolve.py @@ -8,10 +8,9 @@ from functools import partial from typing import Any, Literal, Optional -from numpy.typing import ArrayLike, NDArray - import jax import jax.numpy as jnp +from numpy.typing import ArrayLike, NDArray from . import type_casting, utils diff --git a/src/nemos/glm.py b/src/nemos/glm.py index 70b4ee9e..44cbb4d4 100644 --- a/src/nemos/glm.py +++ b/src/nemos/glm.py @@ -7,13 +7,12 @@ from functools import wraps from typing import Any, Callable, Literal, NamedTuple, Optional, Tuple, Union +import jax +import jax.numpy as jnp import jaxopt from numpy.typing import ArrayLike from scipy.optimize import root -import jax -import jax.numpy as jnp - from . import observation_models as obs from . import tree_utils, validation from .base_regressor import BaseRegressor diff --git a/src/nemos/observation_models.py b/src/nemos/observation_models.py index de0b8981..9d683ae1 100644 --- a/src/nemos/observation_models.py +++ b/src/nemos/observation_models.py @@ -3,10 +3,9 @@ import abc from typing import Callable, Literal, Union -from numpy.typing import NDArray - import jax import jax.numpy as jnp +from numpy.typing import NDArray from . import utils from .base_class import Base diff --git a/src/nemos/regularizer.py b/src/nemos/regularizer.py index c2cc75d6..6d6cf0bd 100644 --- a/src/nemos/regularizer.py +++ b/src/nemos/regularizer.py @@ -9,11 +9,10 @@ import abc from typing import Callable, Tuple, Union -import jaxopt -from numpy.typing import NDArray - import jax import jax.numpy as jnp +import jaxopt +from numpy.typing import NDArray from . import tree_utils from .base_class import Base diff --git a/src/nemos/simulation.py b/src/nemos/simulation.py index 92c64c12..e698e702 100644 --- a/src/nemos/simulation.py +++ b/src/nemos/simulation.py @@ -2,13 +2,12 @@ from typing import Callable, Tuple, Union +import jax +import jax.numpy as jnp import numpy as np import scipy.stats as sts from numpy.typing import NDArray -import jax -import jax.numpy as jnp - from . import convolve, validation from .pytrees import FeaturePytree diff --git a/src/nemos/solvers.py b/src/nemos/solvers.py index 779d24d0..4c060609 100644 --- a/src/nemos/solvers.py +++ b/src/nemos/solvers.py @@ -1,14 +1,13 @@ from functools import partial from typing import Callable, NamedTuple, Optional, Union -from jaxopt import OptStep -from jaxopt._src import loop -from jaxopt.prox import prox_none - import jax import jax.flatten_util import jax.numpy as jnp from jax import grad, jit, lax, random +from jaxopt import OptStep +from jaxopt._src import loop +from jaxopt.prox import prox_none from .tree_utils import tree_add_scalar_mul, tree_l2_norm, tree_slice, tree_sub from .typing import KeyArrayLike, Pytree diff --git a/src/nemos/type_casting.py b/src/nemos/type_casting.py index 02a73dd8..8a5522b8 100644 --- a/src/nemos/type_casting.py +++ b/src/nemos/type_casting.py @@ -12,13 +12,12 @@ from functools import wraps from typing import Any, Callable, List, Literal, Optional, Type, Union +import jax +import jax.numpy as jnp import numpy as np import pynapple as nap from numpy.typing import NDArray -import jax -import jax.numpy as jnp - from . import tree_utils _NAP_TIME_PRECISION = 10 ** (-nap.nap_config.time_index_precision) diff --git a/src/nemos/typing.py b/src/nemos/typing.py index fa86ca82..dd9bc5a6 100644 --- a/src/nemos/typing.py +++ b/src/nemos/typing.py @@ -2,9 +2,8 @@ from typing import Any, Callable, NamedTuple, Tuple, Union -import jaxopt - import jax.numpy as jnp +import jaxopt from jax._src.typing import ArrayLike from .pytrees import FeaturePytree diff --git a/src/nemos/utils.py b/src/nemos/utils.py index f457e4f1..87ad0472 100644 --- a/src/nemos/utils.py +++ b/src/nemos/utils.py @@ -3,11 +3,10 @@ import warnings from typing import Any, Callable, List, Literal, Union -import numpy as np -from numpy.typing import NDArray - import jax import jax.numpy as jnp +import numpy as np +from numpy.typing import NDArray from .tree_utils import pytree_map_and_reduce from .type_casting import is_numpy_array_like, support_pynapple diff --git a/src/nemos/validation.py b/src/nemos/validation.py index c3ff63bf..9dd0853f 100644 --- a/src/nemos/validation.py +++ b/src/nemos/validation.py @@ -3,10 +3,9 @@ import warnings from typing import Any, Optional, Union -from numpy.typing import DTypeLike, NDArray - import jax import jax.numpy as jnp +from numpy.typing import DTypeLike, NDArray from .pytrees import FeaturePytree from .tree_utils import get_valid_multitree, pytree_map_and_reduce