tony/tensorneat-mend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add input_transform and update_input_transform;

Browse Source 
change the args for genome.forward.
Origin: (state, inputs, transformed)
New: (state, transformed, inputs)
This commit is contained in:
wls2002
2024-06-03 10:53:15 +08:00
parent a07a3b1cb2
commit edfb0596e7
16 changed files with 185 additions and 221 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tensorneat/problem/func_fit/func_fit.py
View File

@@ -20,8 +20,8 @@ class FuncFit(BaseProblem):
def evaluate(self, state, randkey, act_func, params):
predict = jax.vmap(act_func, in_axes=(None, 0, None))(
state, self.inputs, params
predict = jax.vmap(act_func, in_axes=(None, None, 0))(
state, params, self.inputs
)
if self.error_method == "mse":
@@ -45,8 +45,8 @@ class FuncFit(BaseProblem):
return -loss
def show(self, state, randkey, act_func, params, *args, **kwargs):
predict = jax.vmap(act_func, in_axes=(None, 0, None))(
state, self.inputs, params
predict = jax.vmap(act_func, in_axes=(None, None, 0))(
state, params, self.inputs, params
)
inputs, target, predict = jax.device_get([self.inputs, self.targets, predict])
if self.return_data: