Add CustomFuncFit into problem; Add related examples

tensorneat/common/__init__.py

@@ -31,6 +31,7 @@ name2sympy = {
     "maxabs": SympyMaxabs,
     "mean": SympyMean,
     "clip": SympyClip,
+    "square": SympySquare,
 }
 
 

tensorneat/common/activation/act_jnp.py

@@ -69,6 +69,10 @@ class Act:
         z = jnp.clip(z, -10, 10)
         return jnp.exp(z)
 
+    @staticmethod
+    def square(z):
+        return jnp.pow(z, 2)
+
     @staticmethod
     def abs(z):
         z = jnp.clip(z, -1, 1)

tensorneat/common/activation/act_sympy.py

@@ -184,6 +184,12 @@ class SympyExp(sp.Function):
         return rf"\mathrm{{exp}}\left({sp.latex(self.args[0])}\right)"
 
 
+class SympySquare(sp.Function):
+    @classmethod
+    def eval(cls, z):
+        return sp.Pow(z, 2)
+
+
 class SympyAbs(sp.Function):
     @classmethod
     def eval(cls, z):