moscot 0.4.2__tar.gz → 0.5.0__tar.gz

{moscot-0.4.2 → moscot-0.5.0}/.pre-commit-config.yaml

@@ -63,7 +63,7 @@ repos:
       - id: doc8
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.9.10
+    rev: v0.11.7
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{moscot-0.4.2 → moscot-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: moscot
-Version: 0.4.2
+Version: 0.5.0
 Summary: Multi-omic single-cell optimal transport tools
 Author: Dominik Klein, Giovanni Palla, Michal Klein, Zoe Piran, Marius Lange
 Maintainer-email: Dominik Klein <dominik.klein@helmholtz-muenchen.de>, Giovanni Palla <giovanni.palla@helmholtz-muenchen.de>, Michal Klein <michal.klein@helmholtz-muenchen.de>
@@ -65,18 +65,18 @@ Requires-Dist: anndata>=0.9.1
 Requires-Dist: scanpy>=1.9.3
 Requires-Dist: wrapt>=1.13.2
 Requires-Dist: docrep>=0.3.2
-Requires-Dist: ott-jax>=0.5.0
+Requires-Dist: jax>=0.6.1
+Requires-Dist: ott-jax>=0.6.0
 Requires-Dist: cloudpickle>=2.2.0
 Requires-Dist: rich>=13.5
 Requires-Dist: docstring_inheritance>=2.0.0
 Requires-Dist: mudata>=0.2.2
+Requires-Dist: optax
+Requires-Dist: flax
+Requires-Dist: diffrax
+Requires-Dist: ott-jax[neural]>=0.5.0
 Provides-Extra: spatial
 Requires-Dist: squidpy>=1.2.3; extra == "spatial"
-Provides-Extra: neural
-Requires-Dist: optax; extra == "neural"
-Requires-Dist: flax; extra == "neural"
-Requires-Dist: diffrax; extra == "neural"
-Requires-Dist: ott-jax[neural]>=0.5.0; extra == "neural"
 Provides-Extra: dev
 Requires-Dist: pre-commit>=3.0.0; extra == "dev"
 Requires-Dist: tox>=4; extra == "dev"

{moscot-0.4.2 → moscot-0.5.0}/pyproject.toml

@@ -54,11 +54,16 @@ dependencies = [
     "scanpy>=1.9.3",
     "wrapt>=1.13.2",
     "docrep>=0.3.2",
-    "ott-jax>=0.5.0",
+    "jax>=0.6.1",
+    "ott-jax>=0.6.0",
     "cloudpickle>=2.2.0",
     "rich>=13.5",
     "docstring_inheritance>=2.0.0",
-    "mudata>=0.2.2"
+    "mudata>=0.2.2",
+    "optax",
+    "flax",
+    "diffrax",
+    "ott-jax[neural]>=0.5.0"
 ]
 
 [project.optional-dependencies]
@@ -66,13 +71,6 @@ spatial = [
     "squidpy>=1.2.3"
 ]
 
-neural = [
-    "optax",
-    "flax",
-    "diffrax",
-    "ott-jax[neural]>=0.5.0",
-
-]
 dev = [
     "pre-commit>=3.0.0",
     "tox>=4",

{moscot-0.4.2 → moscot-0.5.0}/src/moscot/backends/ott/__init__.py

@@ -1,15 +1,14 @@
 from ott.geometry import costs
 
 from moscot.backends.ott._utils import sinkhorn_divergence
-from moscot.backends.ott.output import GraphOTTOutput, NeuralOutput, OTTOutput
-from moscot.backends.ott.solver import GENOTLinSolver, GWSolver, SinkhornSolver
+from moscot.backends.ott.output import GraphOTTOutput, OTTOutput
+from moscot.backends.ott.solver import GWSolver, SinkhornSolver
 from moscot.costs import register_cost
 
 __all__ = [
     "OTTOutput",
     "GWSolver",
     "SinkhornSolver",
-    "NeuralOutput",
     "sinkhorn_divergence",
     "GENOTLinSolver",
     "GraphOTTOutput",

{moscot-0.4.2 → moscot-0.5.0}/src/moscot/backends/ott/_utils.py

@@ -10,7 +10,6 @@ import scipy.sparse as sp
 from ott.geometry import epsilon_scheduler, geodesic, geometry, pointcloud
 from ott.initializers.linear import initializers as init_lib
 from ott.initializers.linear import initializers_lr as lr_init_lib
-from ott.neural import datasets
 from ott.solvers import utils as solver_utils
 from ott.tools.sinkhorn_divergence import sinkhorn_divergence as sinkhorn_div
 
@@ -18,6 +17,7 @@ from moscot._logging import logger
 from moscot._types import ArrayLike, ScaleCost_t
 
 Scale_t = Union[float, Literal["mean", "median", "max_cost", "max_norm", "max_bound"]]
+OTDataset = Any  # to be removed when neural part is being removed from moscot
 
 
 __all__ = ["sinkhorn_divergence"]
@@ -272,7 +272,7 @@ def data_match_fn(
 
 class Loader:
 
-    def __init__(self, dataset: datasets.OTDataset, batch_size: int, seed: Optional[int] = None):
+    def __init__(self, dataset: OTDataset, batch_size: int, seed: Optional[int] = None):
         self.dataset = dataset
         self.batch_size = batch_size
         self._rng = np.random.default_rng(seed)

{moscot-0.4.2 → moscot-0.5.0}/src/moscot/backends/ott/output.py

@@ -1,12 +1,8 @@
-from typing import Any, Callable, List, Optional, Tuple, Union
-
-import jaxlib.xla_extension as xla_ext
+from typing import Any, Optional, Tuple, Union
 
 import jax
 import jax.numpy as jnp
 import numpy as np
-import scipy.sparse as sp
-from ott.neural.methods.flows.genot import GENOT
 from ott.solvers.linear import sinkhorn, sinkhorn_lr
 from ott.solvers.quadratic import gromov_wasserstein, gromov_wasserstein_lr
 
@@ -14,10 +10,9 @@ import matplotlib as mpl
 import matplotlib.pyplot as plt
 
 from moscot._types import ArrayLike, Device_t
-from moscot.backends.ott._utils import get_nearest_neighbors
-from moscot.base.output import BaseDiscreteSolverOutput, BaseNeuralOutput
+from moscot.base.output import BaseDiscreteSolverOutput
 
-__all__ = ["OTTOutput", "GraphOTTOutput", "NeuralOutput"]
+__all__ = ["OTTOutput", "GraphOTTOutput"]
 
 
 class OTTOutput(BaseDiscreteSolverOutput):
@@ -209,7 +204,7 @@ class OTTOutput(BaseDiscreteSolverOutput):
         else:
             idx = 0
 
-        if not isinstance(device, xla_ext.Device):
+        if not isinstance(device, jax.Device):
             try:
                 device = jax.devices(device)[idx]
             except IndexError:
@@ -244,220 +239,6 @@ class OTTOutput(BaseDiscreteSolverOutput):
         return jnp.ones((n,))
 
 
-class NeuralOutput(BaseNeuralOutput):
-    """Output wrapper for GENOT."""
-
-    def __init__(self, model: GENOT, logs: dict[str, list[float]]):
-        """Initialize `NeuralOutput`.
-
-        Parameters
-        ----------
-        model : GENOT
-            The OTT-Jax GENOT model
-        """
-        self._logs = logs
-        self._model = model
-
-    @property
-    def logs(self):
-        """Logs of the training.  A dictionary containing what the numeric values are i.e., loss.
-
-        Returns
-        -------
-        dict[str, list[float]]
-        """
-        return self._logs
-
-    def _project_transport_matrix(
-        self,
-        src_dist: ArrayLike,
-        tgt_dist: ArrayLike,
-        func: Callable[[ArrayLike], ArrayLike],
-        save_transport_matrix: bool = False,  # TODO(@MUCDK) adapt order of arguments
-        batch_size: int = 1024,
-        k: int = 30,
-        length_scale: Optional[float] = None,
-        seed: int = 42,
-        recall_target: float = 0.95,
-        aggregate_to_topk: bool = True,
-    ) -> sp.csr_matrix:
-        row_indices: List[ArrayLike] = []
-        column_indices: List[ArrayLike] = []
-        distances_list: List[ArrayLike] = []
-        if length_scale is None:
-            key = jax.random.PRNGKey(seed)
-            src_batch = src_dist[jax.random.choice(key, src_dist.shape[0], shape=((batch_size,)))]
-            tgt_batch = tgt_dist[jax.random.choice(key, tgt_dist.shape[0], shape=((batch_size,)))]
-            length_scale = jnp.std(jnp.concatenate((func(src_batch), tgt_batch)))
-        for index in range(0, len(src_dist), batch_size):
-            distances, indices = get_nearest_neighbors(
-                func(src_dist[index : index + batch_size, :]),
-                tgt_dist,
-                k,
-                recall_target=recall_target,
-                aggregate_to_topk=aggregate_to_topk,
-            )
-            distances = jnp.exp(-((distances / length_scale) ** 2))
-            distances /= jnp.expand_dims(jnp.sum(distances, axis=1), axis=1)
-            distances_list.append(distances.flatten())
-            column_indices.append(indices.flatten())
-            row_indices.append(
-                jnp.repeat(jnp.arange(index, index + min(batch_size, len(src_dist) - index)), min(k, len(tgt_dist)))
-            )
-        distances = jnp.concatenate(distances_list)
-        row_indices = jnp.concatenate(row_indices)
-        column_indices = jnp.concatenate(column_indices)
-        tm = sp.csr_matrix((distances, (row_indices, column_indices)), shape=[len(src_dist), len(tgt_dist)])
-        if save_transport_matrix:
-            self._transport_matrix = tm
-        return tm
-
-    def project_to_transport_matrix(  # type:ignore[override]
-        self,
-        src_cells: ArrayLike,
-        tgt_cells: ArrayLike,
-        condition: ArrayLike = None,
-        save_transport_matrix: bool = False,  # TODO(@MUCDK) adapt order of arguments
-        batch_size: int = 1024,
-        k: int = 30,
-        length_scale: Optional[float] = None,
-        seed: int = 42,
-        recall_target: float = 0.95,
-        aggregate_to_topk: bool = True,
-    ) -> sp.csr_matrix:
-        """Project conditional neural OT map onto cells.
-
-        In constrast to discrete OT, (conditional) neural OT does not necessarily map cells onto cells,
-        but a cell can also be mapped to a location between two cells. This function computes
-        a pseudo-transport matrix considering the neighborhood of where a cell is mapped to.
-        Therefore, a neighborhood graph of `k` target cells is computed around each transported cell
-        of the source distribution. The assignment likelihood of each mapped cell to the target cells is then
-        computed with a Gaussian kernel with parameter `length_scale`.
-
-        Parameters
-        ----------
-        condition
-            Condition `src_cells` correspond to.
-        src_cells
-            Cells which are to be mapped.
-        tgt_cells
-            Cells from which the neighborhood graph around the mapped `src_cells` are computed.
-        forward
-            Whether to map cells based on the forward transport map or backward transport map.
-        save_transport_matrix
-            Whether to save the transport matrix.
-        batch_size
-            Number of data points in the source distribution the neighborhood graph is computed
-            for in parallel.
-        k
-            Number of neighbors to construct the k-nearest neighbor graph of a mapped cell.
-        length_scale
-            Length scale of the Gaussian kernel used to compute the assignment likelihood. If `None`,
-            `length_scale` is set to the empirical standard deviation of `batch_size` pairs of data points of the
-            mapped source and target distribution.
-        seed
-            Random seed for sampling the pairs of distributions for computing the variance in case `length_scale`
-            is `None`.
-        recall_target
-            Recall target for the approximation.
-        aggregate_to_topk
-            When true, the nearest neighbor aggregates approximate results to the top-k in sorted order.
-            When false, returns the approximate results unsorted.
-            In this case, the number of the approximate results is implementation defined and is greater or
-            equal to the specified k.
-
-        Returns
-        -------
-        The projected transport matrix.
-        """
-        src_cells, tgt_cells = jnp.asarray(src_cells), jnp.asarray(tgt_cells)
-        conditioned_fn: Callable[[ArrayLike], ArrayLike] = lambda x: self.push(x, condition)
-        push = self.push if condition is None else conditioned_fn
-        func, src_dist, tgt_dist = (push, src_cells, tgt_cells)
-        return self._project_transport_matrix(
-            src_dist=src_dist,
-            tgt_dist=tgt_dist,
-            func=func,
-            save_transport_matrix=save_transport_matrix,  # TODO(@MUCDK) adapt order of arguments
-            batch_size=batch_size,
-            k=k,
-            length_scale=length_scale,
-            seed=seed,
-            recall_target=recall_target,
-            aggregate_to_topk=aggregate_to_topk,
-        )
-
-    def push(self, x: ArrayLike, cond: Optional[ArrayLike] = None) -> ArrayLike:
-        """Push distribution `x` conditioned on condition `cond`.
-
-        Parameters
-        ----------
-        x
-            Distribution to push.
-        cond
-            Condition of conditional neural OT.
-
-        Returns
-        -------
-        Pushed distribution.
-        """
-        if isinstance(x, (bool, int, float, complex)):
-            raise ValueError("Expected array, found scalar value.")
-        if x.ndim not in (1, 2):
-            raise ValueError(f"Expected 1D or 2D array, found `{x.ndim}`.")
-        return self._apply_forward(x, cond=cond)
-
-    def _apply_forward(self, x: ArrayLike, cond: Optional[ArrayLike] = None) -> ArrayLike:
-        return self._model.transport(x, condition=cond)
-
-    @property
-    def is_linear(self) -> bool:  # noqa: D102
-        return True  # TODO(ilan-gold): need to contribute something to ott-jax so this is resolvable from GENOT
-
-    @property
-    def shape(self) -> Tuple[int, int]:
-        """%(shape)s."""
-        raise NotImplementedError()
-
-    def to(
-        self,
-        device: Optional[Device_t] = None,
-    ) -> "NeuralOutput":
-        """Transfer the output to another device or change its data type.
-
-        Parameters
-        ----------
-        device
-            If not `None`, the output will be transferred to `device`.
-
-        Returns
-        -------
-        The output on a saved on `device`.
-        """
-        # # TODO(michalk8): when polishing docs, move the definition to the base class + use docrep
-        # if isinstance(device, str) and ":" in device:
-        #     device, ix = device.split(":")
-        #     idx = int(ix)
-        # else:
-        #     idx = 0
-
-        # if not isinstance(device, xla_ext.Device):
-        #     try:
-        #         device = jax.devices(device)[idx]
-        #     except IndexError as err:
-        #         raise IndexError(f"Unable to fetch the device with `id={idx}`.") from err
-
-        # out = jax.device_put(self._model, device)
-        # return NeuralOutput(out)
-        return self  # TODO(ilan-gold) move model to device
-
-    @property
-    def converged(self) -> bool:
-        """%(converged)s."""
-        # always return True for now
-        return True
-
-
 class GraphOTTOutput(OTTOutput):
     """Output of :term:`OT` problems with a graph geometry in the linear term.
 
@@ -509,7 +290,7 @@ class GraphOTTOutput(OTTOutput):
         else:
             idx = 0
 
-        if not isinstance(device, xla_ext.Device):
+        if not isinstance(device, jax.Device):
             try:
                 device = jax.devices(device)[idx]
             except IndexError:

{moscot-0.4.2 → moscot-0.5.0}/src/moscot/backends/ott/solver.py

@@ -1,12 +1,9 @@
 import abc
-import functools
 import inspect
-import math
 import types
 from typing import (
     Any,
     Hashable,
-    List,
     Literal,
     Mapping,
     NamedTuple,
@@ -17,21 +14,13 @@ from typing import (
     Union,
 )
 
-import optax
-
 import jax
 import jax.numpy as jnp
-import numpy as np
 from ott.geometry import costs, epsilon_scheduler, geodesic, geometry, pointcloud
-from ott.neural.datasets import OTData, OTDataset
-from ott.neural.methods.flows import dynamics, genot
-from ott.neural.networks.layers import time_encoder
-from ott.neural.networks.velocity_field import VelocityField
 from ott.problems.linear import linear_problem
 from ott.problems.quadratic import quadratic_problem
 from ott.solvers.linear import sinkhorn, sinkhorn_lr
 from ott.solvers.quadratic import gromov_wasserstein, gromov_wasserstein_lr
-from ott.solvers.utils import uniform_sampler
 
 from moscot._logging import logger
 from moscot._types import (
@@ -43,23 +32,20 @@ from moscot._types import (
 )
 from moscot.backends.ott._utils import (
     InitializerResolver,
-    Loader,
-    MultiLoader,
     _instantiate_geodesic_cost,
     alpha_to_fused_penalty,
     check_shapes,
     convert_scipy_sparse,
-    data_match_fn,
     densify,
     ensure_2d,
 )
-from moscot.backends.ott.output import GraphOTTOutput, NeuralOutput, OTTOutput
+from moscot.backends.ott.output import GraphOTTOutput, OTTOutput
 from moscot.base.problems._utils import TimeScalesHeatKernel
 from moscot.base.solver import OTSolver
 from moscot.costs import get_cost
-from moscot.utils.tagged_array import DistributionCollection, TaggedArray
+from moscot.utils.tagged_array import TaggedArray
 
-__all__ = ["SinkhornSolver", "GWSolver", "GENOTLinSolver"]
+__all__ = ["SinkhornSolver", "GWSolver"]
 
 OTTSolver_t = Union[
     sinkhorn.Sinkhorn,
@@ -516,193 +502,3 @@ class GWSolver(OTTJaxSolver):
         problem_kwargs -= {"geom_xx", "geom_yy", "geom_xy", "fused_penalty"}
         problem_kwargs |= {"alpha"}
         return geom_kwargs | problem_kwargs, {"epsilon"}
-
-
-class GENOTLinSolver(OTSolver[OTTOutput]):
-    """Solver class for genot.GENOT linear :cite:`klein2023generative`."""
-
-    def __init__(self, **kwargs: Any) -> None:
-        """Initiate the class with any kwargs passed to the ott-jax class."""
-        super().__init__()
-        self._train_sampler: Optional[MultiLoader] = None
-        self._valid_sampler: Optional[MultiLoader] = None
-        self._neural_kwargs = kwargs
-
-    @property
-    def problem_kind(self) -> ProblemKind_t:  # noqa: D102
-        return "linear"
-
-    def _prepare(  # type: ignore[override]
-        self,
-        distributions: DistributionCollection[K],
-        sample_pairs: List[Tuple[Any, Any]],
-        train_size: float = 0.9,
-        batch_size: int = 1024,
-        is_conditional: bool = True,
-        **kwargs: Any,
-    ) -> Tuple[MultiLoader, MultiLoader]:
-        train_loaders = []
-        validate_loaders = []
-        seed = kwargs.get("seed")
-        is_aligned = kwargs.get("is_aligned", False)
-        if train_size == 1.0:
-            for sample_pair in sample_pairs:
-                source_key = sample_pair[0]
-                target_key = sample_pair[1]
-                src_data = OTData(
-                    lin=distributions[source_key].xy,
-                    condition=distributions[source_key].conditions if is_conditional else None,
-                )
-                tgt_data = OTData(
-                    lin=distributions[target_key].xy,
-                    condition=distributions[target_key].conditions if is_conditional else None,
-                )
-                dataset = OTDataset(src_data=src_data, tgt_data=tgt_data, seed=seed, is_aligned=is_aligned)
-                loader = Loader(dataset, batch_size=batch_size, seed=seed)
-                train_loaders.append(loader)
-                validate_loaders.append(loader)
-        else:
-            if train_size > 1.0 or train_size <= 0.0:
-                raise ValueError("Invalid train_size. Must be: 0 < train_size <= 1")
-
-            seed = kwargs.get("seed", 0)
-            for sample_pair in sample_pairs:
-                source_key = sample_pair[0]
-                target_key = sample_pair[1]
-                source_data: ArrayLike = distributions[source_key].xy
-                target_data: ArrayLike = distributions[target_key].xy
-                source_split_data = self._split_data(
-                    source_data,
-                    conditions=distributions[source_key].conditions,
-                    train_size=train_size,
-                    seed=seed,
-                    a=distributions[source_key].a,
-                    b=distributions[source_key].b,
-                )
-                target_split_data = self._split_data(
-                    target_data,
-                    conditions=distributions[target_key].conditions,
-                    train_size=train_size,
-                    seed=seed,
-                    a=distributions[target_key].a,
-                    b=distributions[target_key].b,
-                )
-                src_data_train = OTData(
-                    lin=source_split_data.data_train,
-                    condition=source_split_data.conditions_train if is_conditional else None,
-                )
-                tgt_data_train = OTData(
-                    lin=target_split_data.data_train,
-                    condition=target_split_data.conditions_train if is_conditional else None,
-                )
-                train_dataset = OTDataset(
-                    src_data=src_data_train, tgt_data=tgt_data_train, seed=seed, is_aligned=is_aligned
-                )
-                train_loader = Loader(train_dataset, batch_size=batch_size, seed=seed)
-                src_data_validate = OTData(
-                    lin=source_split_data.data_valid,
-                    condition=source_split_data.conditions_valid if is_conditional else None,
-                )
-                tgt_data_validate = OTData(
-                    lin=target_split_data.data_valid,
-                    condition=target_split_data.conditions_valid if is_conditional else None,
-                )
-                validate_dataset = OTDataset(
-                    src_data=src_data_validate, tgt_data=tgt_data_validate, seed=seed, is_aligned=is_aligned
-                )
-                validate_loader = Loader(validate_dataset, batch_size=batch_size, seed=seed)
-                train_loaders.append(train_loader)
-                validate_loaders.append(validate_loader)
-        source_dim = self._neural_kwargs.get("input_dim", 0)
-        target_dim = source_dim
-        condition_dim = self._neural_kwargs.get("cond_dim", 0)
-        # TODO(ilan-gold): What are reasonable defaults here?
-        neural_vf = VelocityField(
-            output_dims=[*self._neural_kwargs.get("velocity_field_output_dims", []), target_dim],
-            condition_dims=(
-                self._neural_kwargs.get("velocity_field_condition_dims", [source_dim + condition_dim])
-                if is_conditional
-                else None
-            ),
-            hidden_dims=self._neural_kwargs.get("velocity_field_hidden_dims", [1024, 1024, 1024]),
-            time_dims=self._neural_kwargs.get("velocity_field_time_dims", None),
-            time_encoder=self._neural_kwargs.get(
-                "velocity_field_time_encoder", functools.partial(time_encoder.cyclical_time_encoder, n_freqs=1024)
-            ),
-        )
-        seed = self._neural_kwargs.get("seed", 0)
-        rng = jax.random.PRNGKey(seed)
-        data_match_fn_kwargs = self._neural_kwargs.get(
-            "data_match_fn_kwargs",
-            {} if "data_match_fn" in self._neural_kwargs else {"epsilon": 1e-1, "tau_a": 1.0, "tau_b": 1.0},
-        )
-        time_sampler = self._neural_kwargs.get("time_sampler", uniform_sampler)
-        optimizer = self._neural_kwargs.get("optimizer", optax.adam(learning_rate=1e-4))
-        self._solver = genot.GENOT(
-            vf=neural_vf,
-            flow=self._neural_kwargs.get(
-                "flow",
-                dynamics.ConstantNoiseFlow(0.1),
-            ),
-            data_match_fn=functools.partial(
-                self._neural_kwargs.get("data_match_fn", data_match_fn), typ="lin", **data_match_fn_kwargs
-            ),
-            source_dim=source_dim,
-            target_dim=target_dim,
-            condition_dim=condition_dim if is_conditional else None,
-            optimizer=optimizer,
-            time_sampler=time_sampler,
-            rng=rng,
-            latent_noise_fn=self._neural_kwargs.get("latent_noise_fn", None),
-            **self._neural_kwargs.get("velocity_field_train_state_kwargs", {}),
-        )
-        return (
-            MultiLoader(datasets=train_loaders, seed=seed),
-            MultiLoader(datasets=validate_loaders, seed=seed),
-        )
-
-    def _split_data(  # TODO: adapt for Gromov terms
-        self,
-        x: ArrayLike,
-        conditions: Optional[ArrayLike],
-        train_size: float,
-        seed: int,
-        a: Optional[ArrayLike] = None,
-        b: Optional[ArrayLike] = None,
-    ) -> SingleDistributionData:
-        n_samples_x = x.shape[0]
-        n_train_x = math.ceil(train_size * n_samples_x)
-        rng = np.random.default_rng(seed)
-        x = rng.permutation(x)
-        if a is not None:
-            a = rng.permutation(a)
-        if b is not None:
-            b = rng.permutation(b)
-
-        return SingleDistributionData(
-            data_train=x[:n_train_x],
-            data_valid=x[n_train_x:],
-            conditions_train=conditions[:n_train_x] if conditions is not None else None,
-            conditions_valid=conditions[n_train_x:] if conditions is not None else None,
-            a_train=a[:n_train_x] if a is not None else None,
-            a_valid=a[n_train_x:] if a is not None else None,
-            b_train=b[:n_train_x] if b is not None else None,
-            b_valid=b[n_train_x:] if b is not None else None,
-        )
-
-    @property
-    def solver(self) -> genot.GENOT:
-        """Underlying optimal transport solver."""
-        return self._solver
-
-    @classmethod
-    def _call_kwargs(cls) -> Tuple[Set[str], Set[str]]:
-        return {"batch_size", "train_size", "trainloader", "validloader", "seed"}, {}  # type: ignore[return-value]
-
-    def _solve(self, data_samplers: Tuple[MultiLoader, MultiLoader]) -> NeuralOutput:  # type: ignore[override]
-        seed = self._neural_kwargs.get("seed", 0)  # TODO(ilan-gold): unify rng hadnling like OTT tests
-        rng = jax.random.PRNGKey(seed)
-        logs = self.solver(
-            data_samplers[0], n_iters=self._neural_kwargs.get("n_iters", 100), rng=rng
-        )  # TODO(ilan-gold): validation and figure out defualts
-        return NeuralOutput(self.solver, logs)