- u = torch.randint(self.len_total, (nb,))
- v = torch.randint(self.len_total, (nb,))
-
- a = torch.randint(self.nb_values, (nb,))
- b = torch.randint(self.nb_values, (nb,))
- c = torch.randint(self.nb_values, (nb,))
-
- while True:
- to_compute = torch.logical_or(u>=v-self.len_total//10,u<v-self.len_total//5)
- to_compute =torch.logical_or(to_compute, u == 0)
- to_compute =torch.logical_or(to_compute, v == self.len_total)
- n = to_compute.long().sum()
- if n == 0:
- break
- else:
- u[to_compute] = torch.randint(self.len_total, (n,))
- v[to_compute] = torch.randint(self.len_total, (n,))
-
- while True:
- to_compute = a==b
- to_compute = torch.logical_or(to_compute,b==c)
- to_compute = torch.logical_or(to_compute,a==c)
-
- if self.global_constraint:
- to_compute = torch.logical_or(to_compute,(a*u+b*(v-u)+c*(self.len_total-v)) // self.len_total != self.nb_values//2)
-
- n = to_compute.long().sum()
- if n == 0:
- break
- else:
- a[to_compute] = torch.randint(self.nb_values, (n,))
- b[to_compute] = torch.randint(self.nb_values, (n,))
- c[to_compute] = torch.randint(self.nb_values, (n,))
-
- assert (u>=v).long().sum() == 0
- assert (a==b).long().sum() == 0
- assert (a==c).long().sum() == 0
- assert (c==b).long().sum() == 0
-
- t = torch.arange(self.len_total)
- seq = (t[None,:] < u[:,None]).long() * a[:,None] + \
- (t[None,:] >= u[:,None]).long() * (t[None,:] < v[:,None]).long() * b[:,None] + \
- (t[None,:] >= v[:,None]).long() * c[:,None]
+ x = (torch.rand(nb,self.nb_state_tokens).sort(dim=-1).indices == 0).long() * self.value_max
+ seq = [x]
+
+ for t in range(self.nb_time_steps-1):
+ v = torch.rand(x.size()) * (x > 0).float()
+ u = (v.max(dim=-1,keepdim=True).values == v).long()
+ n = (u*x*torch.rand(x.size())).long().sum(dim=-1,keepdim=True) // 2
+ x = x + n * (u.roll(shifts=-1,dims=-1) - 2 * u + u.roll(shifts=1,dims=-1))
+ seq.append(x)