python-ldl 0.0.3__py3-none-any.whl → 0.1.0__py3-none-any.whl

pyldl/algorithms/__init__.py

@@ -1,30 +1,198 @@
-from ._problem_transformation import PT_Bayes, PT_SVM, LDSVR
-from ._algorithm_adaptation import AA_BP, AA_KNN, CAD, QFD2, CJS, CPNN, BCPNN, ACPNN
-from ._specialized_algorithms import SA_BFGS, SA_IIS
+from typing import TYPE_CHECKING
 
-from ._incomplete import IncomLDL
-from ._classifier import LDL4C, LDL_HR, LDLM
-from ._ensemble import DF_LDL, AdaBoostLDL
+if TYPE_CHECKING:
+    from ._problem_transformation import *
+    from ._algorithm_adaptation import *
+    from ._specialized_algorithms import *
 
-from ._ldlf import LDLF
-from ._ldl_scl import LDL_SCL
-from ._ldl_lrr import LDL_LRR
-from ._ldl_dpa import LDL_DPA
+    from ._incomplete import *
+    from ._classifier import *
+    from ._ensemble import *
 
-from ._ssg_ldl import SSG_LDL
+    from ._bp import *
+    from ._cpnn import *
+    from ._duo_ldl import *
+    from ._ldlf import *
+    from ._ldllc import *
+    from ._ldlsf import *
+    from ._ldl_lclr import *
+    from ._ldl_scl import *
+    from ._ldl_lrr import *
+    from ._ldl_dpa import *
+    from ._lrldl import *
+    from ._ldl_hvlc import *
+    from ._rknn_ldl import *
+    from ._s_ldl import *
+    from ._delta_ldl import *
+    from ._snefy_ldl import *
+    from ._ldl_dvs import *
+    from ._ldl_dpm import *
 
-from ._label_enhancement import FCM, KM, LP, ML, GLLE, LEVI, LIBLE
+    from ._ssg_ldl import *
 
-from ._ldl_da import LDL_DA
+    from ._label_enhancement import *
 
+    from ._ldl_da import *
 
-__all__ = ["SA_BFGS", "SA_IIS", "AA_KNN", "AA_BP", "PT_Bayes", "PT_SVM",
-           "CPNN", "BCPNN", "ACPNN", "LDSVR",
-           "LDLF", "LDL_SCL", "LDL_LRR", "LDL_DPA", "CAD", "QFD2", "CJS",
-           "DF_LDL", "AdaBoostLDL",
-           "LDL4C", "LDL_HR", "LDLM",
-           "IncomLDL",
-           "SSG_LDL",
-           "FCM", "KM", "LP", "ML", "GLLE",
-           "LEVI", "LIBLE",
-           "LDL_DA"]
+
+_LDL_MODULE_MAP = {
+    # -------------------- 2026 --------------------
+    "LDL_DVS": "._ldl_dvs",
+    "LDL_DPM": "._ldl_dpm",
+    # -------------------- 2025 --------------------
+    "RG4LDL": "._ensemble",
+    "RKNN_LDL": "._rknn_ldl",
+    "SNEFY_LDL": "._snefy_ldl",
+    **dict.fromkeys(
+        [
+            "_S_LDL",
+            "S_LRR",
+            "S_SCL",
+            "S_KLD",
+            "S_CJS",
+            "S_QFD2",
+            "Shallow_S_LDL",
+        ],
+        "._s_ldl"
+    ),
+    "Delta_LDL": "._delta_ldl",
+# -------------------- 2024 --------------------
+    "LDL_HVLC": "._ldl_hvlc",
+    **dict.fromkeys(
+        [
+            "_LRLDL",
+            "TKLRLDL",
+            "TLRLDL",
+        ],
+        "._lrldl"
+    ),
+    "LDL_DPA": "._ldl_dpa",
+    # -------------------- 2023 --------------------
+    "LDL_LRR": "._ldl_lrr",
+    # -------------------- 2021 --------------------
+    "DF_LDL": "._ensemble",
+    "LDL_SCL": "._ldl_scl",
+    "Duo_LDL": "._duo_ldl",
+    "BD_LDL": "._algorithm_adaptation",
+    # -------------------- 2019 --------------------
+    "LDL_LCLR": "._ldl_lclr",
+    "LDLSF": "._ldlsf",
+    # -------------------- 2018 --------------------
+    "LDLLC": "._ldllc",
+    "LALOT": "._specialized_algorithms",
+    "StructRF": "._ensemble",
+    # -------------------- 2017 --------------------
+    "BCPNN": "._cpnn",
+    "ACPNN": "._cpnn",
+    "LDLF": "._ldlf",
+    # -------------------- 2016 --------------------
+    "LDLogitBoost": "._ensemble",
+    **dict.fromkeys(
+        [
+            "_SA",
+            "SA_BFGS",
+            "SA_IIS",
+        ],
+        "._specialized_algorithms"
+    ),
+    "AA_KNN": "._algorithm_adaptation",
+    "AA_BP": "._bp",
+    **dict.fromkeys(
+        [
+            "_PT",
+            "PT_Bayes",
+            "PT_SVM",
+        ],
+        "._problem_transformation"
+    ),
+    # -------------------- 2015 --------------------
+    "LDSVR": "._problem_transformation",
+    # -------------------- 2013 --------------------
+    "CPNN": "._cpnn",
+}
+
+_LE_MODULE_MAP = {
+    # -------------------- 2023 --------------------
+    **dict.fromkeys(
+        [
+            "LIBLE",
+            "ConLE",
+        ],
+        "._label_enhancement"
+    ),
+    # -------------------- 2020 --------------------
+    "LEVI": "._label_enhancement",
+    # -------------------- 2019 --------------------
+    **dict.fromkeys(
+        [
+            "GLLE",
+            "ML",
+            "LP",
+            "KM",
+            "FCM",
+        ],
+        "._label_enhancement"
+    ),
+}
+
+_INCOMLDL_MODULE_MAP = {
+    # -------------------- 2024 --------------------
+    "WInLDL": "._incomplete",
+    # -------------------- 2017 --------------------
+    "IncomLDL": "._incomplete",
+}
+
+_LDL4C_MODULE_MAP = {
+    # -------------------- 2021 --------------------
+    **dict.fromkeys(
+        [
+            "LDLM",
+            "LDL_HR",
+            "LDL4C",
+        ],
+        "._classifier"
+    ),
+}
+
+_SSG_LDL_MODULE_MAP = {
+    # -------------------- 2021 --------------------
+    "SSG_LDL": "._ssg_ldl"
+}
+
+_LDL_DA_MODULE_MAP = {
+    # -------------------- 2025 --------------------
+    "LDL_DA": "._ldl_da"
+}
+
+_GLD_MODULE_MAP = {
+    # -------------------- 2026 --------------------
+    "GLD_SVR": "._problem_transformation",
+    "GLD_KNN": "._algorithm_adaptation",  
+    "GLD_BFGS": "._specialized_algorithms",
+}
+
+_MODULE_MAP = {
+    **_LDL_MODULE_MAP,
+    **_LE_MODULE_MAP,
+    **_INCOMLDL_MODULE_MAP,
+    **_LDL4C_MODULE_MAP,
+    **_SSG_LDL_MODULE_MAP,
+    **_LDL_DA_MODULE_MAP,
+    **_GLD_MODULE_MAP,
+}
+
+
+_ldl__ = list(_LDL_MODULE_MAP)
+_le__ = list(_LE_MODULE_MAP)
+_incomldl__ = list(_INCOMLDL_MODULE_MAP)
+_ldl4c__ = list(_LDL4C_MODULE_MAP)
+_ssg_ldl__ = list(_SSG_LDL_MODULE_MAP)
+_ldl_da__ = list(_LDL_DA_MODULE_MAP)
+_gld__ = list(_GLD_MODULE_MAP)
+
+
+__all__ = _ldl__ + _le__ + _incomldl__ + _ldl4c__ + _ssg_ldl__ + _ldl_da__ + _gld__
+
+
+from .base._lazy import lazy_module
+__getattr__ = lazy_module(_MODULE_MAP, __name__)

pyldl/algorithms/_algorithm_adaptation.py

@@ -1,198 +1,63 @@
 import numpy as np
-from sklearn.neighbors import NearestNeighbors
 
-import keras
-import tensorflow as tf
-from keras import backend as K
+from pyldl.algorithms.base import Base, BaseLDL, BaseGLD
 
-from pyldl.algorithms.base import BaseLDL, BaseDeepLDL, BaseGD, BaseAdam
 
+EPS = np.finfo(np.float32).eps
 
-class AA_KNN(BaseLDL):
 
-    def fit(self, X, y, k=5):
-        super().fit(X, y)
-        self._model = NearestNeighbors(n_neighbors=k).fit(self._X)
+class _AA_KNN(Base):
+    def __init__(self, *, k: int = 5, **kwargs):
+        super().__init__(**kwargs)
+        self.k = k
+
+    def fit(self, X: np.ndarray, target: np.ndarray):
+        from sklearn.neighbors import NearestNeighbors
+        super().fit(X, target)
+        self._knn = NearestNeighbors(n_neighbors=self.k).fit(self._X)
         return self
 
     def predict(self, X):
-        _, inds = self._model.kneighbors(X)
-        return np.average(self._y[inds], axis=1)
+        _, inds = self._knn.kneighbors(X)
+        return np.average(self._target[inds], axis=1)
+
+    def __getstate__(self):
+        return self.__dict__.copy()
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+
+
+class AA_KNN(_AA_KNN, BaseLDL):
+    """:class:`AA-kNN <pyldl.algorithms.AA_KNN>` is proposed in paper :cite:`2016:geng`.
+    """
+    pass
 
 
-class AA_BP(BaseGD, BaseDeepLDL):
+class GLD_KNN(_AA_KNN, BaseGLD):
+    """:class:`GLD-kNN <pyldl.algorithms.GLD_KNN>` is proposed in paper :cite:`2026:wu`.
+    """
     pass
 
 
-class CAD(BaseAdam, BaseDeepLDL):
-
-    @staticmethod
-    @tf.function
-    def loss_function(y, y_pred):
-        def _CAD(y, y_pred):
-            return tf.reduce_mean(tf.abs(
-                tf.cumsum(y, axis=1) - tf.cumsum(y_pred, axis=1)
-            ), axis=1)
-        return tf.math.reduce_sum(
-            tf.map_fn(lambda i: _CAD(y[:, :i], y_pred[:, :i]),
-                      tf.range(1, y.shape[1] + 1),
-                      fn_output_signature=tf.float32)
-        )
-
-
-class QFD2(BaseAdam, BaseDeepLDL):
-
-    @staticmethod
-    @tf.function
-    def _loss_function(y, y_pred):
-        Q = y - y_pred
-        j = tf.reshape(tf.range(y.shape[1]), [y.shape[1], 1])
-        k = tf.reshape(tf.range(y.shape[1]), [1, y.shape[1]])
-        A = tf.cast(1 - tf.abs(j - k) / (y.shape[1] - 1), dtype=tf.float32)
-        return tf.math.reduce_mean(
-            tf.linalg.diag_part(tf.matmul(tf.matmul(Q, A), tf.transpose(Q)))
-        )
-
-
-class CJS(BaseAdam, BaseDeepLDL):
-
-    @staticmethod
-    @tf.function
-    def loss_function(y, y_pred):
-        def _CJS(y, y_pred):
-            m = 0.5 * (y + y_pred)
-            js = 0.5 * (keras.losses.kl_divergence(y, m) + keras.losses.kl_divergence(y_pred, m))
-            return tf.reduce_mean(js)
-        return tf.math.reduce_sum(
-            tf.map_fn(lambda i: _CJS(y[:, :i], y_pred[:, :i]),
-                      tf.range(1, y.shape[1] + 1),
-                      fn_output_signature=tf.float32)
-        )
-
-
-class RProp(keras.optimizers.Optimizer):
-    
-    def __init__(self, init_alpha=1e-3, scale_up=1.2, scale_down=0.5, min_alpha=1e-6, max_alpha=50., **kwargs):
-        super(RProp, self).__init__(name='rprop', **kwargs)
-        self.init_alpha = K.variable(init_alpha, name='init_alpha')
-        self.scale_up = K.variable(scale_up, name='scale_up')
-        self.scale_down = K.variable(scale_down, name='scale_down')
-        self.min_alpha = K.variable(min_alpha, name='min_alpha')
-        self.max_alpha = K.variable(max_alpha, name='max_alpha')
-
-    def apply_gradients(self, grads_and_vars):
-        grads, trainable_variables = zip(*grads_and_vars)
-        self.get_updates(trainable_variables, grads)
-
-    def get_updates(self, params, gradients):
-        grads = gradients
-        shapes = [K.int_shape(p) for p in params]
-        alphas = [K.variable(np.ones(shape) * self.init_alpha) for shape in shapes]
-        old_grads = [K.zeros(shape) for shape in shapes]
-        prev_weight_deltas = [K.zeros(shape) for shape in shapes]
-        self.updates = []
-
-        for param, grad, old_grad, prev_weight_delta, alpha in zip(params, grads,
-                                                                   old_grads, prev_weight_deltas,
-                                                                   alphas):
-
-            new_alpha = K.switch(
-                K.greater(grad * old_grad, 0),
-                K.minimum(alpha * self.scale_up, self.max_alpha),
-                K.switch(K.less(grad * old_grad, 0), K.maximum(alpha * self.scale_down, self.min_alpha), alpha)
-            )
-
-            new_delta = K.switch(K.greater(grad, 0),
-                                 -new_alpha,
-                                 K.switch(K.less(grad, 0),
-                                          new_alpha,
-                                          K.zeros_like(new_alpha)))
-
-            weight_delta = K.switch(K.less(grad*old_grad, 0), -prev_weight_delta, new_delta)
-
-            new_param = param + weight_delta
-
-            grad = K.switch(K.less(grad*old_grad, 0), K.zeros_like(grad), grad)
-
-            self.updates.append(K.update(param, new_param))
-            self.updates.append(K.update(alpha, new_alpha))
-            self.updates.append(K.update(old_grad, grad))
-            self.updates.append(K.update(prev_weight_delta, weight_delta))
-
-        return self.updates
-
-    def get_config(self):
-        config = {
-            'init_alpha': float(K.get_value(self.init_alpha)),
-            'scale_up': float(K.get_value(self.scale_up)),
-            'scale_down': float(K.get_value(self.scale_down)),
-            'min_alpha': float(K.get_value(self.min_alpha)),
-            'max_alpha': float(K.get_value(self.max_alpha)),
-        }
-        base_config = super(RProp, self).get_config()
-        return dict(list(base_config.items()) + list(config.items()))
-
-
-class CPNN(BaseGD, BaseDeepLDL):
-
-    def _not_proper_mode(self):
-        raise ValueError("The argument 'mode' can only be 'none', 'binary' or 'augment'.")
-
-    def __init__(self, mode='none', v=5, n_hidden=64, n_latent=None, random_state=None):
-        super().__init__(n_hidden, n_latent, random_state)
-        if mode == 'none' or mode == 'binary' or mode == 'augment':
-            self._mode = mode
-        else:
-            self._not_proper_mode()
-        self._v = v
-
-    @staticmethod
-    @tf.function
-    def loss_function(y, y_pred):
-        return tf.math.reduce_mean(keras.losses.kl_divergence(y, y_pred))
-
-    def _get_default_model(self):
-        input_shape = (self._n_features + (1 if self._mode == 'none' else self._n_outputs),)
-        return keras.Sequential([keras.layers.InputLayer(input_shape=input_shape),
-                                 keras.layers.Dense(self._n_hidden, activation='sigmoid'),
-                                 keras.layers.Dense(1, activation=None)])
-
-    def _get_default_optimizer(self):
-        return RProp()
-
-    def _before_train(self):
-        if self._mode == 'augment':
-            n = self._X.shape[0]
-            one_hot = tf.one_hot(tf.math.argmax(self._y, axis=1), self.n_outputs)
-            self._X = tf.repeat(self._X, self._v, axis=0)
-            self._y = tf.repeat(self._y, self._v, axis=0)
-            one_hot = tf.repeat(one_hot, self._v, axis=0)
-            v = tf.reshape(tf.tile([1 / (i + 1) for i in range(self._v)], [n]), (-1, 1))
-            self._y += self._y * one_hot * v
-
-    def _make_inputs(self, X):
-        temp = tf.reshape(tf.tile([i + 1 for i in range(self._n_outputs)], [X.shape[0]]), (-1, 1))
-        if self._mode != 'none':
-            temp = tf.one_hot(tf.reshape(temp, (-1, )) - 1, depth=self._n_outputs)
-        return tf.concat([tf.cast(tf.repeat(X, self._n_outputs, axis=0), dtype=tf.float32),
-                          tf.cast(temp, dtype=tf.float32)],
-                          axis=1)
-
-    def _call(self, X):
-        inputs = self._make_inputs(X)
-        outputs = self._model(inputs)
-        results = tf.reshape(outputs, (X.shape[0], self._n_outputs))
-        b = tf.reshape(-tf.math.log(tf.math.reduce_sum(tf.math.exp(results), axis=1)), (-1, 1))
-        return tf.math.exp(b + results)
-
-
-class BCPNN(CPNN):
-
-    def __init__(self, **params):
-        super().__init__(mode='binary', **params)
-
-
-class ACPNN(CPNN):
-
-    def __init__(self, **params):
-        super().__init__(mode='augment', **params)
+class BD_LDL(BaseLDL):
+    """:class:`BD-LDL <pyldl.algorithms.BD_LDL>` is proposed in paper :cite:`2021:liu2`.
+    """
+
+    def __init__(self, *, alpha: float = 1e-3, beta: float = 1e-2, **kwargs):
+        super().__init__(**kwargs)
+        self.alpha = alpha
+        self.beta = beta
+
+    def fit(self, X, D):
+        from scipy.linalg import solve_sylvester
+        super().fit(X, D)
+        A = self._X.T @ self._X + self.beta * np.eye(self._n_features)
+        B = self.alpha * self._D.T @ self._D
+        C = (1 + self.alpha) * self._X.T @ self._D
+        self._W = solve_sylvester(A, B, C)
+        return self
+
+    def predict(self, X):
+        from pyldl.algorithms.utils import proj
+        return proj(X @ self._W)

pyldl/algorithms/_bp.py

@@ -0,0 +1,12 @@
+import keras
+
+from pyldl.algorithms.base import BaseDeepLDL, BaseGD
+
+
+@keras.saving.register_keras_serializable()
+class AA_BP(BaseGD, BaseDeepLDL):
+    """:class:`AA-BP <pyldl.algorithms.AA_BP>` is proposed in paper :cite:`2016:geng`.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)

pyldl/algorithms/_classifier.py

@@ -1,105 +1,131 @@
+import numpy as np
+
 import keras
-import tensorflow as tf
+import keras.ops as ops
 
 from pyldl.algorithms.base import BaseDeepLDLClassifier, BaseGD, BaseBFGS
 
 
-class LDL4C(BaseBFGS, BaseDeepLDLClassifier):
-    """LDL4C is proposed in paper Classification with Label Distribution Learning.
-    """
+EPS = np.finfo(np.float32).eps
 
-    @tf.function
-    def _loss(self, params_1d):
-        y_pred = keras.activations.softmax(self._X @ self._params2model(params_1d)[0])
-        top2 = tf.gather(y_pred, self._top2, axis=1, batch_dims=1)
-        margin = tf.reduce_sum(tf.maximum(0., 1. - (top2[:, 0] - top2[:, 1]) / self._rho))
-        mae = keras.losses.mean_absolute_error(self._y, y_pred)
-        return tf.reduce_sum(self._entropy * mae) + self._alpha * margin + self._beta * self._l2_reg(self._model)
 
-    def _before_train(self):
-        self._top2 = tf.math.top_k(self._y, k=2)[1]
-        self._entropy = tf.cast(-tf.reduce_sum(self._y * tf.math.log(self._y) + 1e-7, axis=1), dtype=tf.float32)
+class LDL4C(BaseBFGS, BaseDeepLDLClassifier):
+    """:class:`LDL4C <pyldl.algorithms.LDL4C>` is proposed in paper :cite:`2019:wang3`.
+    """
 
-    def fit(self, X, y, alpha=1e-2, beta=1e-6, rho=1e-2, **kwargs):
+    def __init__(self, alpha=1e-2, beta=0., rho=1e-2, **kwargs):
+        super().__init__(**kwargs)
         self._alpha = alpha
         self._beta = beta
         self._rho = rho
-        return super().fit(X, y, **kwargs)
+
+    def _loss(self, params_1d):
+        theta = self._params2model(params_1d)[0]
+        D_pred = keras.activations.softmax(self._X @ theta)
+        top2 = ops.take_along_axis(D_pred, self._top2, axis=1)
+        margin = ops.sum(
+            ops.maximum(0., 1. - (top2[:, 0] - top2[:, 1]) / self._rho)
+        )
+        mae = keras.losses.mean_absolute_error(self._D, D_pred)
+        return ops.sum(self._entropy * mae) + self._alpha * margin + self._beta * self._l2_reg(theta)
+
+    def _before_train(self):
+        _, self._top2 = ops.top_k(self._D, k=2)
+        self._entropy = ops.cast(
+            -ops.sum(self._D * ops.log(self._D + EPS), axis=1),
+            dtype="float32")
 
 
 class LDL_HR(BaseBFGS, BaseDeepLDLClassifier):
-    """LDL-HR is proposed in paper Learn the Highest Label and Rest Label Description Degrees.
+    """:class:`LDL-HR <pyldl.algorithms.LDL_HR>` is proposed in paper :cite:`2021:wang3`.
     """
 
-    @tf.function
+    def __init__(self, alpha=1e-2, beta=1e-2, gamma=0., rho=1e-2, **kwargs):
+        super().__init__(**kwargs)
+        self._alpha = alpha
+        self._beta = beta
+        self._gamma = gamma
+        self._rho = rho
+
     def _loss(self, params_1d):
-        y_pred = keras.activations.softmax(self._X @ self._params2model(params_1d)[0])
+        theta = self._params2model(params_1d)[0]
+        D_pred = keras.activations.softmax(self._X @ theta)
 
-        highest = tf.gather(y_pred, self._highest, axis=1, batch_dims=1)
-        rest = tf.gather(y_pred, self._rest, axis=1, batch_dims=1)
-        margin = tf.reduce_sum(tf.maximum(0., 1. - (highest - rest) / self._rho))
+        highest = ops.take_along_axis(D_pred, self._highest, axis=1)
+        rest = ops.take_along_axis(D_pred, self._rest, axis=1)
+        margin = ops.sum(
+            ops.maximum(0., 1. - (highest - rest) / self._rho)
+        )
 
-        real_rest = tf.gather(self._y, self._rest, axis=1, batch_dims=1)
-        rest_mae = tf.reduce_sum(keras.losses.mean_absolute_error(real_rest, rest))
+        real_rest = ops.take_along_axis(self._D, self._rest, axis=1)
+        rest_mae = ops.sum(
+            keras.losses.mean_absolute_error(real_rest, rest)
+        )
 
-        mae = tf.reduce_sum(keras.losses.mean_absolute_error(self._l, y_pred))
+        mae = ops.sum(keras.losses.mean_absolute_error(self._L, D_pred))
 
-        return mae + self._alpha * margin + self._beta * rest_mae + self._gamma * self._l2_reg(self._model)
+        return mae + self._alpha * margin + self._beta * rest_mae + self._gamma * self._l2_reg(theta)
 
     def _before_train(self):
-        temp = tf.math.top_k(self._y, k=self._n_outputs)[1]
+        _, temp = ops.top_k(self._D, k=self._n_outputs)
         self._highest = temp[:, 0:1]
         self._rest = temp[:, 1:]
-        self._l = tf.one_hot(tf.reshape(self._highest, -1), self._n_outputs)
-
-    def fit(self, X, y, alpha=1e-2, beta=1e-2, gamma=1e-6, rho=1e-2, **kwargs):
-        self._alpha = alpha
-        self._beta = beta
-        self._gamma = gamma
-        self._rho = rho
-        return super().fit(X, y, **kwargs) 
+        self._L = ops.one_hot(ops.reshape(self._highest, (-1,)), self._n_outputs)
 
 
 class LDLM(BaseGD, BaseDeepLDLClassifier):
-    """LDLM is proposed in paper Label Distribution Learning Machine.
+    """:class:`LDLM <pyldl.algorithms.LDLM>` is proposed in paper :cite:`2021:wang`.
     """
 
-    @tf.function
-    def _loss(self, X, y, start, end):
-        y_pred = self._model(X)
-
-        pred_margin = tf.reduce_sum(tf.clip_by_value(
-            tf.reduce_sum(tf.abs(self._l[start:end] - y_pred), axis=1) - self._rho,
-            0., float('inf')))
-
-        highest = tf.gather(y_pred, self._highest[start:end], axis=1, batch_dims=1)
-        rest = tf.gather(y_pred, self._rest[start:end], axis=1, batch_dims=1)
-        label_margin = tf.reduce_sum(tf.maximum(0., 1. - (highest - rest) / self._rho))
+    def __init__(self, alpha=1e-2, beta=1e-2, gamma=0., rho=1e-2, **kwargs):
+        super().__init__(**kwargs)
+        self._alpha = alpha
+        self._beta = beta
+        self._gamma = gamma
+        self._rho = rho
 
-        second_margin = tf.reduce_sum(tf.clip_by_value(
-            tf.reduce_sum(tf.abs((y - y_pred) * self._neg_l[start:end]), axis=1) - self._second_margin[start:end],
-            0., float('inf')))
+    def _loss(self, X, D, start, end):
+        D_pred = self._model(X)
+
+        pred_margin = ops.sum(
+            ops.clip(
+                ops.sum(
+                    ops.abs(self._L[start:end] - D_pred), axis=1
+                ) - self._rho, 0., float('inf')
+            )
+        )
+
+        highest = ops.take_along_axis(
+            D_pred, self._highest[start:end], axis=1
+        )
+        rest = ops.take_along_axis(D_pred, self._rest[start:end], axis=1)
+        label_margin = ops.sum(
+            ops.maximum(0., 1. - (highest - rest) / self._rho)
+        )
+
+        second_margin = ops.sum(
+            ops.clip(
+                ops.sum(
+                    ops.abs((D - D_pred) * self._neg_L[start:end]), axis=1
+                ) - self._second_margin[start:end], 0., float('inf')
+            )
+        )
 
         return pred_margin + self._alpha * label_margin + \
-            self._beta * second_margin + self._gamma * self._l2_reg(self._model)
+            self._beta * second_margin + \
+            self._gamma * self._l2_reg(self._model)
 
     def _before_train(self):
-        temp = tf.math.top_k(self._y, k=self._n_outputs)[1]
+        _, temp = ops.top_k(self._D, k=self._n_outputs)
         self._highest = temp[:, 0:1]
         self._rest = temp[:, 1:]
 
-        temp = tf.math.top_k(tf.gather(self._y, self._rest, axis=1, batch_dims=1), k=2)[0]
+        rest = ops.take_along_axis(self._D, self._rest, axis=1)
+        temp, _ = ops.top_k(rest, k=2)
         self._second_margin = temp[:, 0] - temp[:, 1]
 
-        self._l = tf.one_hot(tf.reshape(self._highest, -1), self._n_outputs)
-        self._neg_l = tf.where(tf.equal(self._l, 0.), 1., 0.)
+        self._L = ops.one_hot(ops.reshape(self._highest, (-1,)), self._n_outputs)
+        self._neg_L = ops.where(ops.equal(self._L, 0.), 1., 0.)
 
     def _get_default_model(self):
         return self.get_2layer_model(self._n_features, self._n_outputs)
-
-    def fit(self, X, y, alpha=1e-2, beta=1e-2, gamma=1e-6, rho=1e-2, **kwargs):
-        self._alpha = alpha
-        self._beta = beta
-        self._gamma = gamma
-        self._rho = rho
-        return super().fit(X, y, **kwargs)