5bd6e5c357
add flatten_conns as an inverse function for unflatten_conns; add "test_flatten.ipynb" as test for them.
39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
import jax, jax.numpy as jnp
|
|
from .. import BaseGene
|
|
|
|
|
|
class BaseConnGene(BaseGene):
|
|
"Base class for connection genes."
|
|
fixed_attrs = ["input_index", "output_index", "enabled"]
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def crossover(self, state, randkey, gene1, gene2):
|
|
def crossover_attr():
|
|
return jnp.where(
|
|
jax.random.normal(randkey, gene1.shape) > 0,
|
|
gene1,
|
|
gene2,
|
|
)
|
|
|
|
return jax.lax.cond(
|
|
gene1[2] == gene2[2], # if both genes are enabled or disabled
|
|
crossover_attr, # then randomly pick attributes from gene1 or gene2
|
|
lambda: jnp.where( # one gene is enabled and the other is disabled
|
|
gene1[2], # if gene1 is enabled
|
|
gene1, # then return gene1
|
|
gene2, # else return gene2
|
|
),
|
|
)
|
|
|
|
def forward(self, state, attrs, inputs):
|
|
raise NotImplementedError
|
|
|
|
def update_by_batch(self, state, attrs, batch_inputs):
|
|
# default: do not update attrs, but to calculate batch_res
|
|
return (
|
|
jax.vmap(self.forward, in_axes=(None, None, 0))(state, attrs, batch_inputs),
|
|
attrs,
|
|
)
|