def execute(batch_size, latent_size, data_size):
data = np.random.normal(size=(data_size,
latent_size)).astype(np.float32)
def step_state(state):
return state + np.sum(np.tensordot(data, state, ([1], [1])))
state = np.random.normal(size=(batch_size,
latent_size)).astype(np.float32)
def choose_depth(count):
del count
return 3
program = test_programs.pea_nuts_program((latent_size, ),
choose_depth, step_state)
input_counts = np.array([3] * batch_size)
return vm.execute(program, [input_counts, state],
10,
backend=NP_BACKEND)
def execute(_, backend):
data = tf.random.normal(shape=(data_size, latent_size),
dtype=np.float32)
def step_state(state):
return state + tf.reduce_sum(
input_tensor=tf.tensordot(data, state, ([1], [1])))
state = tf.random.normal(shape=(batch_size, latent_size),
dtype=np.float32)
def choose_depth(count):
del count
return 2
program = test_programs.pea_nuts_program((latent_size, ),
choose_depth, step_state)
input_counts = np.array([3] * batch_size)
return vm.execute(program, [input_counts, state],
10,
backend=backend)