refactor RnnModel: multi-unit delay training and unified vc interface

- Add batch_delay() to shift visible input per unit index
- Unify forward_step() to work with combined vc matrix
- Fix split() to always slice on axis=1
- Add index param to Entity for readable naming
- Rename test_xor.py to xor.py, replace Mat with np.array

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

JayRnn.py

@@ -24,10 +24,10 @@ from rbm.train import train
 from rbm.status import Status
 
 # ── Hyper-parameters ──────────────────────────────────────────────────────────
-TEXT    = "HALLO SUPER JENS UND SUPER MAUSI!" # the character sequence to learn (matches diagram example)
+TEXT    = "0123456789" # the character sequence to learn (matches diagram example)
 WIN    = 3         # sliding-window width  (= N in the diagram)
 STRIDE = 1     	   # sliding-window step size
-UNITS  = 1
+UNITS  = 3
 H_SIZE = 32         # hidden units per RBM cell
 NUM_EPOCHS = 1000
 NUM_ITERATIONS = 1
@@ -67,29 +67,48 @@ def to_batch(_win_str: list[str]):
 		_batch[i, :] = str2vec(win).flatten()
 	return _batch
 
+def batch_delay(_batch: Mat, delay=0):
+	result = _batch
+	if delay > 0:
+		result[0:-delay] = _batch[delay:]
+		result[-delay:] = np.zeros([delay, _batch.shape[1]])
+	return result
+
+def vc2char(_vc: Mat):
+	_v, _ = split(_vc.reshape(1, WIN*vocab_size()+H_SIZE), H_SIZE, axis=1)
+	return vec2str(_v.reshape([WIN, vocab_size()]), axis=1)
 
 class RnnModel(Model):
 	def __init__(self, name: str, work_dir: str = '.'):
 		super().__init__(name, work_dir)
 
 		self.units: list[Entity] = []
-		for _ in range(UNITS):
+		for index in range(UNITS):
-			unit = Entity((WIN*vocab_size() + H_SIZE, H_SIZE), EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False), training_params=TRAIN_PARAMS)
+			unit = Entity((WIN*vocab_size() + H_SIZE, H_SIZE), EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False), training_params=TRAIN_PARAMS, index=index)
 			self.units.append(unit)
 
 	def train(self, vc: Mat, status: Status = None):
-		for unit in self.units:
+		_v, _c = split(vc, H_SIZE, axis=1)
-			train(unit, vc, status)
+		for delay, unit in enumerate(self.units):
+			_vd = batch_delay(_v, delay)
+			_vc = concat(_vd, _c, axis=1)
+			for i in range(vc.shape[0]):
+				print(f"train: {unit.name}:{vc2char(_vc[i,:])}")
+			train(unit, _vc, status)
 			# For the next unit: Update context portion of vc
-			_c = unit.forward(vc)
+			_c = unit.forward(_vc)
-			vc[1:, WIN*vocab_size():] = _c[0:-1,:]
 
-	def forward_step(self, _v: Mat, _c: Mat) -> tuple[Mat, Mat]:
+
-		_vc = concat(_v.reshape([1, WIN*vocab_size()]), _c.reshape([1, H_SIZE]), axis=1)
+	def forward_step(self, _vc: Mat):
 		for unit in self.units:
 			_c = unit.forward(_vc)
 			_vc = unit.reconstruct(_c)
-		return split(_vc.flatten(), H_SIZE)
+			_v, _ = split(_vc, H_SIZE, axis=1)
+			_v_cl = clamp(_v.reshape([WIN, vocab_size()]), axis=1).reshape([1, WIN * vocab_size()])
+			_v = shift_left(_v_cl, 1)
+			_vc = concat(_v, _c, axis=1)
+			print(f"forward_step: {unit.name}:{vc2char(_vc)}")
+		return _vc
 
 	def forward(self, x: Mat):
 		pass
@@ -118,26 +137,29 @@ if __name__ == "__main__":
 
 	# vc contains vis + context
 	# context will be updated after training
-	c_train = np.zeros([len(batch), H_SIZE])
+	c_train = np.zeros([batch.shape[0], H_SIZE])
 	vc_train = concat(batch, c_train, axis=1)
-	model.train(vc_train, status=Status())
+#	model.train(vc_train, status=Status())
-	model.save()
+#	model.save()
 
 	# test the model
-	seed_str = 'HA^'
+	seed_str = '1'
-	seed_padded = ' '*(WIN-len(seed_str)) + seed_str
+	seed_padded = '^'*(WIN-len(seed_str)) + seed_str
-	v_forward = str2vec(seed_padded).flatten()
+	v_test = str2vec(seed_padded).reshape([1, WIN*vocab_size()])
-	context = c_train[0]
+	c_test = np.zeros([1, H_SIZE])
-	text_predict = ''
+	vc_test = concat(v_test, c_test, axis=1)
-	for i in range(len(TEXT)+5):
+	for i in range(len(TEXT) + 5):
-		v_mat = v_forward.reshape([WIN, vocab_size()])
+		vc_test = model.forward_step(vc_test)
-		print(f"Forward {i:02d}: {vec2str(v_mat, axis=1)}")
+	if 0:
-		v_predict, context = model.forward_step(v_mat, context)
+		for i in range(len(TEXT)+5):
-		v_predict_mat = v_predict.reshape([WIN, vocab_size()])
+			v_mat = v_forward.reshape([WIN, vocab_size()])
-		v_predict_clamp = clamp(v_predict_mat, axis=1)
+			print(f"Forward {i:02d}: {vec2str(v_mat, axis=1)}")
-		v_predict_str = vec2str(v_predict_clamp, axis=1)
+			v_predict, context = model.forward_step(v_mat, context)
-		text_predict += v_predict_str[WIN-1]
+			v_predict_mat = v_predict.reshape([WIN, vocab_size()])
-		print(f"Predict {i:02d}: {v_predict_str}")
+			v_predict_clamp = clamp(v_predict_mat, axis=1)
-		v_forward = shift_left(v_predict_clamp.reshape([1, WIN*vocab_size()]), vocab_size())
+			v_predict_str = vec2str(v_predict_clamp, axis=1)
+			text_predict += v_predict_str[WIN-1]
+			print(f"Predict {i:02d}: {v_predict_str}")
+			v_forward = shift_left(v_predict_clamp.reshape([1, WIN*vocab_size()]), vocab_size())
 
-	print(seed_padded[0:WIN-1] + text_predict)
+		print(seed_padded[0:WIN-1] + text_predict)

src/rbm/entity.py

@@ -61,14 +61,14 @@ class Entity:
 		GB_RBM = "GB-RBM"
 		GG_RBM = "GG-RBM"
 
-	def __init__(self, shape: tuple[int, int], params: EntityParams, training_params: TrainingParams|None = None, enable_training: bool = True):
+	def __init__(self, shape: tuple[int, int], params: EntityParams, training_params: TrainingParams|None = None, enable_training: bool = True, index=0):
 		self.shape = shape
 		self.params = params
 		self.training_params = training_params
 		self.enable_training = enable_training
 		self.state = RbmState.from_layer_params(shape)
 		self.grad = RbmState.from_layer_params(shape)
-		self.name = f"Entity-{shape[0]}x{shape[1]}"
+		self.name = f"Entity{index}-{shape[0]}x{shape[1]}"
 		self.type = None
 		if params.do_gaussian_visible:
 			if params.do_gaussian_hidden:

src/stack/rnn_helper.py

@@ -37,7 +37,7 @@ def concat(v: np.ndarray, c: np.ndarray, axis=0) -> np.ndarray:
 
 def split(vc: np.ndarray, h_size: int, axis=0) -> tuple[np.ndarray, np.ndarray]:
 	v_len = vc.shape[axis]-h_size
-	return vc[0:v_len], vc[v_len:]
+	return vc[:, 0:v_len], vc[:, v_len:]
 
 def shift_right(m: np.ndarray, amount: int = 1) -> np.ndarray:
 	return np.hstack([np.zeros((m.shape[0], amount)), m[:, :-amount]])

tests/xor.py

@@ -26,7 +26,7 @@ def xor():
 	layer.load(os.path.join(WORK_DIR, "xor_layer0_state.npz"))
 
 	# Prepare training data
-	training_batch = Mat([[0,1,1], [0,0,0], [1,1,0], [1,0,1]], dtype=np.float64)
+	training_batch = np.array([[0,1,1], [0,0,0], [1,1,0], [1,0,1]], dtype=np.float64)
 
 	# Train layer
 	train(layer.entity, training_batch, Status())