add package problems

problems/function_fitting/function_fitting_problem.py

@@ -0,0 +1,22 @@
+import numpy as np
+import jax
+
+from problems import Problem
+
+
+class FunctionFittingProblem(Problem):
+    def __init__(self, num_inputs, num_outputs, batch, inputs, target, loss='MSE'):
+        self.forward_way = 'pop_batch'
+        self.num_inputs = num_inputs
+        self.num_outputs = num_outputs
+        self.batch = batch
+        self.inputs = inputs
+        self.target = target
+        self.loss = loss
+        super().__init__(self.forward_way, self.num_inputs, self.num_outputs, self.batch)
+
+    def evaluate(self, batch_forward_func):
+        out = batch_forward_func(self.inputs)
+        out = jax.device_get(out)
+        fitnesses = 1 - np.mean((self.target - out) ** 2, axis=(1, 2))
+        return fitnesses.tolist()