57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import jax, jax.numpy as jnp
|
|
from .default import DefaultGenome
|
|
|
|
|
|
class DenseInitialize(DefaultGenome):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
assert self.max_nodes >= self.num_inputs + self.num_outputs
|
|
assert self.max_conns >= self.num_inputs * self.num_outputs
|
|
|
|
def initialize(self, state, randkey):
|
|
|
|
k1, k2 = jax.random.split(randkey, num=2)
|
|
|
|
input_idx, output_idx = self.input_idx, self.output_idx
|
|
input_size = len(input_idx)
|
|
output_size = len(output_idx)
|
|
|
|
nodes = jnp.full(
|
|
(self.max_nodes, self.node_gene.length), jnp.nan, dtype=jnp.float32
|
|
)
|
|
|
|
nodes = nodes.at[input_idx, 0].set(input_idx)
|
|
nodes = nodes.at[output_idx, 0].set(output_idx)
|
|
|
|
total_idx = input_size + output_size
|
|
rand_keys_n = jax.random.split(k1, num=total_idx)
|
|
|
|
node_attr_func = jax.vmap(self.node_gene.new_random_attrs, in_axes=(None, 0))
|
|
node_attrs = node_attr_func(state, rand_keys_n)
|
|
nodes = nodes.at[:total_idx, 1:].set(node_attrs)
|
|
|
|
conns = jnp.full(
|
|
(self.max_conns, self.conn_gene.length), jnp.nan, dtype=jnp.float32
|
|
)
|
|
|
|
input_to_output_ids, output_ids = jnp.meshgrid(
|
|
input_idx, output_idx, indexing="ij"
|
|
)
|
|
total_conns = input_size * output_size
|
|
conns = conns.at[:total_conns, :2].set(
|
|
jnp.column_stack([input_to_output_ids.flatten(), output_ids.flatten()])
|
|
)
|
|
|
|
rand_keys_c = jax.random.split(k2, num=total_conns)
|
|
conns_attr_func = jax.vmap(
|
|
self.conn_gene.new_random_attrs,
|
|
in_axes=(
|
|
None,
|
|
0,
|
|
),
|
|
)
|
|
conns_attrs = conns_attr_func(state, rand_keys_c)
|
|
conns = conns.at[:total_conns, 2:].set(conns_attrs)
|
|
|
|
return nodes, conns
|