# Path: brainforge/llatomic/_llactivation.py
# @nb.vectorize(finfout, nopython=True)
# def relu(Z):
#     return s0 if s0 >= Z else Z
# 
# @nb.jit(nopython=True)
# def tanh(Z):
#     return np.tanh(Z)


import numba as nb
import numpy as np
from ._llactivation import relu, tanh


@nb.jit(nopython=True)
def recurrent_forward_relu(X, W, b):
    outdim = W.shape[-1]
    time, batch, indim = X.shape
    O = np.zeros((time, batch, outdim))
    for t in range(time):
        Z = np.concatenate((X[t], O[t-1]), axis=-1)
        preact = np.dot(Z, W) + b

# Path: brainforge/llatomic/_llactivation.py
# @nb.vectorize(finfout, nopython=True)
# def relu(Z):
#     return s0 if s0 >= Z else Z
# 
# @nb.jit(nopython=True)
# def tanh(Z):
#     return np.tanh(Z)


# Path: brainforge/llatomic/_llrecurrent.py
import numba as nb
import numpy as np
from ._llactivation import relu, tanh



@nb.jit(nopython=True)
def recurrent_forward_relu(X, W, b):
    outdim = W.shape[-1]
    time, batch, indim = X.shape
    O = np.zeros((time, batch, outdim))
    for t in range(time):
        Z = np.concatenate((X[t], O[t-1]), axis=-1)
        preact = np.dot(Z, W) + b

        O[t] = relu(preact.ravel()).reshape(*preact.shape)