add backend="jax" to sympy module

2024-06-13 05:55:33 +08:00
parent b3e442c688
commit 69d73aab73
12 changed files with 254 additions and 167 deletions
--- a/tensorneat/utils/aggregation/agg_sympy.py
+++ b/tensorneat/utils/aggregation/agg_sympy.py
@@ -1,3 +1,4 @@
+import numpy as np
 import sympy as sp


@@ -51,15 +52,6 @@ class SympyMedian(sp.Function):

        return None

-    @staticmethod
-    def numerical_eval(args):
-        sorted_args = sorted(args)
-        n = len(sorted_args)
-        if n % 2 == 1:
-            return sorted_args[n // 2]
-        else:
-            return (sorted_args[n // 2 - 1] + sorted_args[n // 2]) / 2
-
    def _sympystr(self, printer):
        return f"median({', '.join(map(str, self.args))})"