edfb0596e7
change the args for genome.forward. Origin: (state, inputs, transformed) New: (state, transformed, inputs)
42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
import jax, jax.numpy as jnp
|
|
from .. import BaseGene
|
|
|
|
|
|
class BaseNodeGene(BaseGene):
|
|
"Base class for node genes."
|
|
fixed_attrs = ["index"]
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, state, attrs, inputs, is_output_node=False):
|
|
raise NotImplementedError
|
|
|
|
def input_transform(self, state, attrs, inputs):
|
|
"""
|
|
make transformation in the input node.
|
|
default: do nothing
|
|
"""
|
|
return inputs
|
|
|
|
def update_by_batch(self, state, attrs, batch_inputs, is_output_node=False):
|
|
# default: do not update attrs, but to calculate batch_res
|
|
return (
|
|
jax.vmap(self.forward, in_axes=(None, None, 0, None))(
|
|
state, attrs, batch_inputs, is_output_node
|
|
),
|
|
attrs,
|
|
)
|
|
|
|
def update_input_transform(self, state, attrs, batch_inputs):
|
|
"""
|
|
update the attrs for transformation in the input node.
|
|
default: do nothing
|
|
"""
|
|
return (
|
|
jax.vmap(self.input_transform, in_axes=(None, None, 0))(
|
|
state, attrs, batch_inputs
|
|
),
|
|
attrs,
|
|
)
|