modifying
This commit is contained in:
@@ -32,14 +32,14 @@ min_species_size = 1
|
||||
|
||||
[gene-bias]
|
||||
bias_init_mean = 0.0
|
||||
bias_init_stdev = 1.0
|
||||
bias_init_std = 1.0
|
||||
bias_mutate_power = 0.5
|
||||
bias_mutate_rate = 0.7
|
||||
bias_replace_rate = 0.1
|
||||
|
||||
[gene-response]
|
||||
response_init_mean = 1.0
|
||||
response_init_stdev = 0.0
|
||||
response_init_std = 0.0
|
||||
response_mutate_power = 0.0
|
||||
response_mutate_rate = 0.0
|
||||
response_replace_rate = 0.0
|
||||
@@ -56,7 +56,7 @@ aggregation_replace_rate = 0.0
|
||||
|
||||
[gene-weight]
|
||||
weight_init_mean = 0.0
|
||||
weight_init_stdev = 1.0
|
||||
weight_init_std = 1.0
|
||||
weight_mutate_power = 0.5
|
||||
weight_mutate_rate = 0.8
|
||||
weight_replace_rate = 0.1
|
||||
|
||||
48
examples/a.py
Normal file
48
examples/a.py
Normal file
@@ -0,0 +1,48 @@
|
||||
import numpy as np
|
||||
|
||||
import jax.numpy as jnp
|
||||
import jax
|
||||
|
||||
a = jnp.array([1, 0, 1, 0, np.nan])
|
||||
b = jnp.array([1, 1, 1, 1, 1])
|
||||
c = jnp.array([1, 1, 1, 1, 1])
|
||||
|
||||
full = jnp.array([
|
||||
[1, 1, 1],
|
||||
[0, 1, 1],
|
||||
[1, 1, 1],
|
||||
[0, 1, 1],
|
||||
])
|
||||
|
||||
print(jnp.column_stack([a[:, None], b[:, None], c[:, None]]))
|
||||
|
||||
aux0 = full[:, 0, None]
|
||||
aux1 = full[:, 1, None]
|
||||
|
||||
print(aux0, aux0.shape)
|
||||
|
||||
print(jnp.concatenate([aux0, aux1], axis=1))
|
||||
|
||||
f_a = jnp.array([False, False, True, True])
|
||||
f_b = jnp.array([True, False, False, False])
|
||||
|
||||
print(jnp.logical_and(f_a, f_b))
|
||||
print(f_a & f_b)
|
||||
|
||||
print(f_a + jnp.nan * 0.0)
|
||||
print(f_a + 1 * 0.0)
|
||||
|
||||
|
||||
@jax.jit
|
||||
def main():
|
||||
return func('happy') + func('sad')
|
||||
|
||||
|
||||
def func(x):
|
||||
if x == 'happy':
|
||||
return 1
|
||||
else:
|
||||
return 2
|
||||
|
||||
|
||||
print(main())
|
||||
@@ -9,6 +9,7 @@ from jax import numpy as jnp
|
||||
|
||||
# from .configs import fetch_first, I_INT
|
||||
from neat.genome.utils import fetch_first, I_INT
|
||||
from .utils import unflatten_connections
|
||||
|
||||
|
||||
@jit
|
||||
@@ -129,6 +130,9 @@ def check_cycles(nodes: Array, connections: Array, from_idx: Array, to_idx: Arra
|
||||
check_cycles(nodes, connections, 0, 3) -> False
|
||||
check_cycles(nodes, connections, 1, 0) -> False
|
||||
"""
|
||||
|
||||
connections = unflatten_connections(nodes, connections)
|
||||
|
||||
connections_enable = ~jnp.isnan(connections[0, :, :])
|
||||
|
||||
connections_enable = connections_enable.at[from_idx, to_idx].set(True)
|
||||
|
||||
@@ -10,7 +10,7 @@ import jax
|
||||
from jax import numpy as jnp
|
||||
from jax import jit, Array
|
||||
|
||||
from .utils import fetch_random, fetch_first, I_INT, unflatten_connections
|
||||
from .utils import fetch_random, fetch_first, I_INT
|
||||
from .genome_ import add_node, delete_node_by_idx, delete_connection_by_idx, add_connection
|
||||
from .graph import check_cycles
|
||||
|
||||
@@ -25,44 +25,30 @@ def mutate(rand_key: Array, nodes: Array, connections: Array, new_node_key: int,
|
||||
:param jit_config:
|
||||
:return:
|
||||
"""
|
||||
|
||||
def m_add_node(rk, n, c):
|
||||
return mutate_add_node(rk, n, c, new_node_key, jit_config['bias_init_mean'], jit_config['response_init_mean'],
|
||||
jit_config['activation_default'], jit_config['aggregation_default'])
|
||||
|
||||
def m_add_connection(rk, n, c):
|
||||
return mutate_add_connection(rk, n, c, jit_config['input_idx'], jit_config['output_idx'])
|
||||
|
||||
def m_delete_node(rk, n, c):
|
||||
return mutate_delete_node(rk, n, c, jit_config['input_idx'], jit_config['output_idx'])
|
||||
|
||||
def m_delete_connection(rk, n, c):
|
||||
return mutate_delete_connection(rk, n, c)
|
||||
|
||||
r1, r2, r3, r4, rand_key = jax.random.split(rand_key, 5)
|
||||
|
||||
# structural mutations
|
||||
# mutate add node
|
||||
r = rand(r1)
|
||||
aux_nodes, aux_connections = m_add_node(r1, nodes, connections)
|
||||
aux_nodes, aux_connections = mutate_add_node(r1, nodes, connections, new_node_key, jit_config)
|
||||
nodes = jnp.where(r < jit_config['node_add_prob'], aux_nodes, nodes)
|
||||
connections = jnp.where(r < jit_config['node_add_prob'], aux_connections, connections)
|
||||
|
||||
# mutate add connection
|
||||
r = rand(r2)
|
||||
aux_nodes, aux_connections = m_add_connection(r3, nodes, connections)
|
||||
aux_nodes, aux_connections = mutate_add_connection(r3, nodes, connections, jit_config)
|
||||
nodes = jnp.where(r < jit_config['conn_add_prob'], aux_nodes, nodes)
|
||||
connections = jnp.where(r < jit_config['conn_add_prob'], aux_connections, connections)
|
||||
|
||||
# mutate delete node
|
||||
r = rand(r3)
|
||||
aux_nodes, aux_connections = m_delete_node(r2, nodes, connections)
|
||||
aux_nodes, aux_connections = mutate_delete_node(r2, nodes, connections, jit_config)
|
||||
nodes = jnp.where(r < jit_config['node_delete_prob'], aux_nodes, nodes)
|
||||
connections = jnp.where(r < jit_config['node_delete_prob'], aux_connections, connections)
|
||||
|
||||
# mutate delete connection
|
||||
r = rand(r4)
|
||||
aux_nodes, aux_connections = m_delete_connection(r4, nodes, connections)
|
||||
aux_nodes, aux_connections = mutate_delete_connection(r4, nodes, connections)
|
||||
nodes = jnp.where(r < jit_config['conn_delete_prob'], aux_nodes, nodes)
|
||||
connections = jnp.where(r < jit_config['conn_delete_prob'], aux_connections, connections)
|
||||
|
||||
@@ -72,7 +58,6 @@ def mutate(rand_key: Array, nodes: Array, connections: Array, new_node_key: int,
|
||||
return nodes, connections
|
||||
|
||||
|
||||
@jit
|
||||
def mutate_values(rand_key: Array, nodes: Array, cons: Array, jit_config: Dict) -> Tuple[Array, Array]:
|
||||
"""
|
||||
Mutate values of nodes and connections.
|
||||
@@ -88,30 +73,41 @@ def mutate_values(rand_key: Array, nodes: Array, cons: Array, jit_config: Dict)
|
||||
"""
|
||||
|
||||
k1, k2, k3, k4, k5, rand_key = jax.random.split(rand_key, num=6)
|
||||
bias_new = mutate_float_values(k1, nodes[:, 1], bias_mean, bias_std,
|
||||
bias_mutate_strength, bias_mutate_rate, bias_replace_rate)
|
||||
response_new = mutate_float_values(k2, nodes[:, 2], response_mean, response_std,
|
||||
response_mutate_strength, response_mutate_rate, response_replace_rate)
|
||||
weight_new = mutate_float_values(k3, cons[:, 2], weight_mean, weight_std,
|
||||
weight_mutate_strength, weight_mutate_rate, weight_replace_rate)
|
||||
act_new = mutate_int_values(k4, nodes[:, 3], act_list, act_replace_rate)
|
||||
agg_new = mutate_int_values(k5, nodes[:, 4], agg_list, agg_replace_rate)
|
||||
|
||||
# mutate enabled
|
||||
# bias
|
||||
bias_new = mutate_float_values(k1, nodes[:, 1], jit_config['bias_init_mean'], jit_config['bias_init_std'],
|
||||
jit_config['bias_mutate_power'], jit_config['bias_mutate_rate'],
|
||||
jit_config['bias_replace_rate'])
|
||||
|
||||
# response
|
||||
response_new = mutate_float_values(k2, nodes[:, 2], jit_config['response_init_mean'],
|
||||
jit_config['response_init_std'], jit_config['response_mutate_power'],
|
||||
jit_config['response_mutate_rate'], jit_config['response_replace_rate'])
|
||||
|
||||
# weight
|
||||
weight_new = mutate_float_values(k3, cons[:, 2], jit_config['weight_init_mean'], jit_config['weight_init_std'],
|
||||
jit_config['weight_mutate_power'], jit_config['weight_mutate_rate'],
|
||||
jit_config['weight_replace_rate'])
|
||||
|
||||
# activation
|
||||
act_new = mutate_int_values(k4, nodes[:, 3], jit_config['activation_options'],
|
||||
jit_config['activation_replace_rate'])
|
||||
|
||||
# aggregation
|
||||
agg_new = mutate_int_values(k5, nodes[:, 4], jit_config['aggregation_options'],
|
||||
jit_config['aggregation_replace_rate'])
|
||||
|
||||
# enabled
|
||||
r = jax.random.uniform(rand_key, cons[:, 3].shape)
|
||||
enabled_new = jnp.where(r < enabled_reverse_rate, 1 - cons[:, 3], cons[:, 3])
|
||||
enabled_new = jnp.where(~jnp.isnan(cons[:, 3]), enabled_new, jnp.nan)
|
||||
enabled_new = jnp.where(r < jit_config['enable_mutate_rate'], 1 - cons[:, 3], cons[:, 3])
|
||||
|
||||
# merge
|
||||
nodes = jnp.column_stack([nodes[:, 0], bias_new, response_new, act_new, agg_new])
|
||||
cons = jnp.column_stack([cons[:, 0], cons[:, 1], weight_new, enabled_new])
|
||||
|
||||
nodes = nodes.at[:, 1].set(bias_new)
|
||||
nodes = nodes.at[:, 2].set(response_new)
|
||||
nodes = nodes.at[:, 3].set(act_new)
|
||||
nodes = nodes.at[:, 4].set(agg_new)
|
||||
cons = cons.at[:, 2].set(weight_new)
|
||||
cons = cons.at[:, 3].set(enabled_new)
|
||||
return nodes, cons
|
||||
|
||||
|
||||
@jit
|
||||
def mutate_float_values(rand_key: Array, old_vals: Array, mean: float, std: float,
|
||||
mutate_strength: float, mutate_rate: float, replace_rate: float) -> Array:
|
||||
"""
|
||||
@@ -132,19 +128,26 @@ def mutate_float_values(rand_key: Array, old_vals: Array, mean: float, std: floa
|
||||
k1, k2, k3, rand_key = jax.random.split(rand_key, num=4)
|
||||
noise = jax.random.normal(k1, old_vals.shape) * mutate_strength
|
||||
replace = jax.random.normal(k2, old_vals.shape) * std + mean
|
||||
|
||||
r = jax.random.uniform(k3, old_vals.shape)
|
||||
|
||||
# default
|
||||
new_vals = old_vals
|
||||
|
||||
# r in [0, mutate_rate), mutate
|
||||
new_vals = jnp.where(r < mutate_rate, new_vals + noise, new_vals)
|
||||
|
||||
# r in [mutate_rate, mutate_rate + replace_rate), replace
|
||||
new_vals = jnp.where(
|
||||
jnp.logical_and(mutate_rate < r, r < mutate_rate + replace_rate),
|
||||
replace,
|
||||
(mutate_rate < r) & (r < mutate_rate + replace_rate),
|
||||
replace + new_vals * 0.0, # in case of nan replace to values
|
||||
new_vals
|
||||
)
|
||||
|
||||
new_vals = jnp.where(~jnp.isnan(old_vals), new_vals, jnp.nan)
|
||||
return new_vals
|
||||
|
||||
|
||||
@jit
|
||||
def mutate_int_values(rand_key: Array, old_vals: Array, val_list: Array, replace_rate: float) -> Array:
|
||||
"""
|
||||
Mutate integer values (act, agg) of a given array.
|
||||
@@ -161,26 +164,20 @@ def mutate_int_values(rand_key: Array, old_vals: Array, val_list: Array, replace
|
||||
k1, k2, rand_key = jax.random.split(rand_key, num=3)
|
||||
replace_val = jax.random.choice(k1, val_list, old_vals.shape)
|
||||
r = jax.random.uniform(k2, old_vals.shape)
|
||||
new_vals = old_vals
|
||||
new_vals = jnp.where(r < replace_rate, replace_val, new_vals)
|
||||
new_vals = jnp.where(~jnp.isnan(old_vals), new_vals, jnp.nan)
|
||||
new_vals = jnp.where(r < replace_rate, replace_val + old_vals * 0.0, old_vals) # in case of nan replace to values
|
||||
|
||||
return new_vals
|
||||
|
||||
|
||||
@jit
|
||||
def mutate_add_node(rand_key: Array, nodes: Array, cons: Array, new_node_key: int,
|
||||
default_bias: float = 0, default_response: float = 1,
|
||||
default_act: int = 0, default_agg: int = 0) -> Tuple[Array, Array]:
|
||||
jit_config: Dict) -> Tuple[Array, Array]:
|
||||
"""
|
||||
Randomly add a new node from splitting a connection.
|
||||
:param rand_key:
|
||||
:param new_node_key:
|
||||
:param nodes:
|
||||
:param cons:
|
||||
:param default_bias:
|
||||
:param default_response:
|
||||
:param default_act:
|
||||
:param default_agg:
|
||||
:param jit_config:
|
||||
:return:
|
||||
"""
|
||||
# randomly choose a connection
|
||||
@@ -192,12 +189,13 @@ def mutate_add_node(rand_key: Array, nodes: Array, cons: Array, new_node_key: in
|
||||
def successful_add_node():
|
||||
# disable the connection
|
||||
new_nodes, new_cons = nodes, cons
|
||||
|
||||
# set enable to false
|
||||
new_cons = new_cons.at[idx, 3].set(False)
|
||||
|
||||
# add a new node
|
||||
new_nodes, new_cons = \
|
||||
add_node(new_nodes, new_cons, new_node_key,
|
||||
bias=default_bias, response=default_response, act=default_act, agg=default_agg)
|
||||
new_nodes, new_cons = add_node(new_nodes, new_cons, new_node_key, bias=0, response=1,
|
||||
act=jit_config['activation_default'], agg=jit_config['aggregation_default'])
|
||||
|
||||
# add two new connections
|
||||
w = new_cons[idx, 2]
|
||||
@@ -211,21 +209,19 @@ def mutate_add_node(rand_key: Array, nodes: Array, cons: Array, new_node_key: in
|
||||
return nodes, cons
|
||||
|
||||
|
||||
# TODO: Need we really need to delete a node?
|
||||
# TODO: Do we really need to delete a node?
|
||||
@jit
|
||||
def mutate_delete_node(rand_key: Array, nodes: Array, cons: Array,
|
||||
input_keys: Array, output_keys: Array) -> Tuple[Array, Array]:
|
||||
def mutate_delete_node(rand_key: Array, nodes: Array, cons: Array, jit_config: Dict) -> Tuple[Array, Array]:
|
||||
"""
|
||||
Randomly delete a node. Input and output nodes are not allowed to be deleted.
|
||||
:param rand_key:
|
||||
:param nodes:
|
||||
:param cons:
|
||||
:param input_keys:
|
||||
:param output_keys:
|
||||
:param jit_config:
|
||||
:return:
|
||||
"""
|
||||
# randomly choose a node
|
||||
node_key, node_idx = choice_node_key(rand_key, nodes, input_keys, output_keys,
|
||||
key, idx = choice_node_key(rand_key, nodes, jit_config['input_idx'], jit_config['output_idx'],
|
||||
allow_input_keys=False, allow_output_keys=False)
|
||||
|
||||
def nothing():
|
||||
@@ -233,37 +229,35 @@ def mutate_delete_node(rand_key: Array, nodes: Array, cons: Array,
|
||||
|
||||
def successful_delete_node():
|
||||
# delete the node
|
||||
aux_nodes, aux_cons = delete_node_by_idx(nodes, cons, node_idx)
|
||||
aux_nodes, aux_cons = delete_node_by_idx(nodes, cons, idx)
|
||||
|
||||
# delete all connections
|
||||
aux_cons = jnp.where(((aux_cons[:, 0] == node_key) | (aux_cons[:, 1] == node_key))[:, jnp.newaxis],
|
||||
aux_cons = jnp.where(((aux_cons[:, 0] == key) | (aux_cons[:, 1] == key))[:, None],
|
||||
jnp.nan, aux_cons)
|
||||
|
||||
return aux_nodes, aux_cons
|
||||
|
||||
nodes, cons = jax.lax.cond(node_idx == I_INT, nothing, successful_delete_node)
|
||||
nodes, cons = jax.lax.cond(idx == I_INT, nothing, successful_delete_node)
|
||||
|
||||
return nodes, cons
|
||||
|
||||
|
||||
@jit
|
||||
def mutate_add_connection(rand_key: Array, nodes: Array, cons: Array,
|
||||
input_keys: Array, output_keys: Array) -> Tuple[Array, Array]:
|
||||
def mutate_add_connection(rand_key: Array, nodes: Array, cons: Array, jit_config: Dict) -> Tuple[Array, Array]:
|
||||
"""
|
||||
Randomly add a new connection. The output node is not allowed to be an input node. If in feedforward networks,
|
||||
cycles are not allowed.
|
||||
:param rand_key:
|
||||
:param nodes:
|
||||
:param cons:
|
||||
:param input_keys:
|
||||
:param output_keys:
|
||||
:param jit_config:
|
||||
:return:
|
||||
"""
|
||||
# randomly choose two nodes
|
||||
k1, k2 = jax.random.split(rand_key, num=2)
|
||||
i_key, from_idx = choice_node_key(k1, nodes, input_keys, output_keys,
|
||||
i_key, from_idx = choice_node_key(k1, nodes, jit_config['input_idx'], jit_config['output_idx'],
|
||||
allow_input_keys=True, allow_output_keys=True)
|
||||
o_key, to_idx = choice_node_key(k2, nodes, input_keys, output_keys,
|
||||
o_key, to_idx = choice_node_key(k2, nodes, jit_config['input_idx'], jit_config['output_idx'],
|
||||
allow_input_keys=False, allow_output_keys=True)
|
||||
|
||||
con_idx = fetch_first((cons[:, 0] == i_key) & (cons[:, 1] == o_key))
|
||||
@@ -280,8 +274,8 @@ def mutate_add_connection(rand_key: Array, nodes: Array, cons: Array,
|
||||
return nodes, cons
|
||||
|
||||
is_already_exist = con_idx != I_INT
|
||||
unflattened = unflatten_connections(nodes, cons)
|
||||
is_cycle = check_cycles(nodes, unflattened, from_idx, to_idx)
|
||||
|
||||
is_cycle = check_cycles(nodes, cons, from_idx, to_idx)
|
||||
|
||||
choice = jnp.where(is_already_exist, 0, jnp.where(is_cycle, 1, 2))
|
||||
nodes, cons = jax.lax.switch(choice, [already_exist, cycle, successful])
|
||||
@@ -311,7 +305,6 @@ def mutate_delete_connection(rand_key: Array, nodes: Array, cons: Array):
|
||||
return nodes, cons
|
||||
|
||||
|
||||
@partial(jit, static_argnames=('allow_input_keys', 'allow_output_keys'))
|
||||
def choice_node_key(rand_key: Array, nodes: Array,
|
||||
input_keys: Array, output_keys: Array,
|
||||
allow_input_keys: bool = False, allow_output_keys: bool = False) -> Tuple[Array, Array]:
|
||||
|
||||
102
neat/genome/utils_.py
Normal file
102
neat/genome/utils_.py
Normal file
@@ -0,0 +1,102 @@
|
||||
from functools import partial
|
||||
|
||||
import jax
|
||||
from jax import numpy as jnp, Array
|
||||
from jax import jit, vmap
|
||||
|
||||
I_INT = jnp.iinfo(jnp.int32).max # infinite int
|
||||
EMPTY_NODE = jnp.full((1, 5), jnp.nan)
|
||||
EMPTY_CON = jnp.full((1, 4), jnp.nan)
|
||||
|
||||
|
||||
@jit
|
||||
def unflatten_connections(nodes: Array, cons: Array):
|
||||
"""
|
||||
transform the (C, 4) connections to (2, N, N)
|
||||
:param nodes: (N, 5)
|
||||
:param cons: (C, 4)
|
||||
:return:
|
||||
"""
|
||||
N = nodes.shape[0]
|
||||
node_keys = nodes[:, 0]
|
||||
i_keys, o_keys = cons[:, 0], cons[:, 1]
|
||||
i_idxs = vmap(fetch_first, in_axes=(0, None))(i_keys, node_keys)
|
||||
i_idxs = key_to_indices(i_keys, node_keys)
|
||||
o_idxs = key_to_indices(o_keys, node_keys)
|
||||
res = jnp.full((2, N, N), jnp.nan)
|
||||
|
||||
# Is interesting that jax use clip when attach data in array
|
||||
# however, it will do nothing set values in an array
|
||||
res = res.at[0, i_idxs, o_idxs].set(cons[:, 2])
|
||||
res = res.at[1, i_idxs, o_idxs].set(cons[:, 3])
|
||||
|
||||
return res
|
||||
|
||||
|
||||
@partial(vmap, in_axes=(0, None))
|
||||
def key_to_indices(key, keys):
|
||||
return fetch_first(key == keys)
|
||||
|
||||
|
||||
@jit
|
||||
def fetch_first(mask, default=I_INT) -> Array:
|
||||
"""
|
||||
fetch the first True index
|
||||
:param mask: array of bool
|
||||
:param default: the default value if no element satisfying the condition
|
||||
:return: the index of the first element satisfying the condition. if no element satisfying the condition, return I_INT
|
||||
example:
|
||||
>>> a = jnp.array([1, 2, 3, 4, 5])
|
||||
>>> fetch_first(a > 3)
|
||||
3
|
||||
>>> fetch_first(a > 30)
|
||||
I_INT
|
||||
"""
|
||||
idx = jnp.argmax(mask)
|
||||
return jnp.where(mask[idx], idx, default)
|
||||
|
||||
|
||||
@jit
|
||||
def fetch_last(mask, default=I_INT) -> Array:
|
||||
"""
|
||||
similar to fetch_first, but fetch the last True index
|
||||
"""
|
||||
reversed_idx = fetch_first(mask[::-1], default)
|
||||
return jnp.where(reversed_idx == -1, -1, mask.shape[0] - reversed_idx - 1)
|
||||
|
||||
|
||||
@jit
|
||||
def fetch_random(rand_key, mask, default=I_INT) -> Array:
|
||||
"""
|
||||
similar to fetch_first, but fetch a random True index
|
||||
"""
|
||||
true_cnt = jnp.sum(mask)
|
||||
cumsum = jnp.cumsum(mask)
|
||||
target = jax.random.randint(rand_key, shape=(), minval=1, maxval=true_cnt + 1)
|
||||
mask = jnp.where(true_cnt == 0, False, cumsum >= target)
|
||||
return fetch_first(mask, default)
|
||||
|
||||
|
||||
@jit
|
||||
def argmin_with_mask(arr: Array, mask: Array) -> Array:
|
||||
masked_arr = jnp.where(mask, arr, jnp.inf)
|
||||
min_idx = jnp.argmin(masked_arr)
|
||||
return min_idx
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
a = jnp.array([1, 2, 3, 4, 5])
|
||||
print(fetch_first(a > 3))
|
||||
print(fetch_first(a > 30))
|
||||
|
||||
print(fetch_last(a > 3))
|
||||
print(fetch_last(a > 30))
|
||||
|
||||
rand_key = jax.random.PRNGKey(0)
|
||||
|
||||
for t in [-1, 0, 1, 2, 3, 4, 5]:
|
||||
for _ in range(10):
|
||||
rand_key, _ = jax.random.split(rand_key)
|
||||
print(jax.random.randint(rand_key, shape=(), minval=1, maxval=2))
|
||||
print(t, fetch_random(rand_key, a > t))
|
||||
Reference in New Issue
Block a user