-
Notifications
You must be signed in to change notification settings - Fork 11
/
Copy pathswap.lua
106 lines (86 loc) · 3.46 KB
/
swap.lua
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
-- This is a config for the Swap task
local config = {}
-- For conversion to distributions
local distUtils = require 'nc.distUtils'
local f = distUtils.flatDist
-- Name of the algorithm
config.name = "Swap"
-- Number of available registers (excluding the RI)
config.nb_registers = 5
-- Number of instructions this program uses
config.nb_existing_ops = 11
-- Size of the memory tape and largest number addressable
config.memory_size = 10
-- Initial state of the registers
config.registers_init = torch.Tensor{0, 1, f, f, 2}
config.registers_init = distUtils.toDistTensor(config.registers_init, config.memory_size)
-- Program
config.nb_states = 9
config.program = {
torch.Tensor{0, 0, 1, 1, 0, 1, 0, 1, f}, -- first arguments
torch.Tensor{f, 4, f, 4, f, f, 3, 2, f}, -- second arguments
torch.Tensor{0, 0, 1, 1, 2, 3, 4, 4, 4}, -- target register
torch.Tensor{8, 3, 8, 3, 8, 8, 9, 9, 0} -- instruction to operate
}
-- Sample input memory
-- the first two values are addresses in the list
-- In this sample, we will switch the third and fourth value (index 2 and 3)
-- Our input list is {4,5,6,7}
-- The output list should be {4,5,7,6}
config.example_input = torch.zeros(config.memory_size)
config.example_input[1] = 2
config.example_input[2] = 3
config.example_input[3] = 4
config.example_input[4] = 5
config.example_input[5] = 6
config.example_input[6] = 7
config.example_output = config.example_input:clone()
config.example_output[5] = 7
config.example_output[6] = 6
config.example_input = distUtils.toDistTensor(config.example_input, config.memory_size)
config.example_output = distUtils.toDistTensor(config.example_output, config.memory_size)
config.example_loss_mask = torch.ones(config.memory_size, config.memory_size)
local function randint(a, b)
return a + math.floor(torch.uniform() * (b-a))
end
config.gen_sample = function()
local possible_list_length = math.floor(config.memory_size-2)
local list_length = randint(2, possible_list_length)
local array_of_value = torch.floor(torch.rand(list_length) * (config.memory_size-1))+1
local p = randint(0, list_length-1)
local q = randint(0, list_length-1)
local input = torch.zeros(config.memory_size)
input[1] = p
input[2] = q
input:narrow(1, 3, list_length):copy(array_of_value)
local output = input:clone()
local loss_mask = torch.zeros(config.memory_size, config.memory_size)
output[3+p] = array_of_value[1+q]
loss_mask[3+p]:fill(1)
output[3+q] = array_of_value[1+p]
loss_mask[3+q]:fill(1)
input = distUtils.toDistTensor(input, config.memory_size)
output = distUtils.toDistTensor(output, config.memory_size)
return input, output, loss_mask
end
config.gen_biased_sample = function()
local possible_list_length = math.floor(config.memory_size-2)
local list_length = 3 + math.floor(torch.uniform() * (possible_list_length-3))
local array_of_value = torch.floor(torch.rand(list_length) * (config.memory_size-1))+1
local p = 0
local q = 2
local input = torch.zeros(config.memory_size)
input[1] = p
input[2] = q
input:narrow(1, 3, list_length):copy(array_of_value)
local output = input:clone()
local loss_mask = torch.zeros(config.memory_size, config.memory_size)
output[3+p] = array_of_value[1+q]
loss_mask[3+p]:fill(1)
output[3+q] = array_of_value[1+p]
loss_mask[3+q]:fill(1)
input = distUtils.toDistTensor(input, config.memory_size)
output = distUtils.toDistTensor(output, config.memory_size)
return input, output, loss_mask
end
return config