--!native
--!optimize 2
local Types = require("../Types")
local Utils = require("./Utils")
local Tensor = require("../Tensor")
local GRU = require("../nn/GRU")
local LSTM = require("../nn/LSTM")
local RNN = require("../nn/RNN")
local Dropout = require("../nn/Dropout")
local LayerNorm = require("../nn/LayerNorm")
local Linear = require("../nn/Linear")

export type SequenceClassifierConfig = {
	inputDim: number,
	hiddenDim: number,
	numLayers: number?,
	numClasses: number,
	cell: "gru" | "lstm" | "rnn"?,
	dropout: number?,
	layerNorm: boolean?,
	finalActivation: Utils.ActivationName?,
}

type CellType = "gru" | "lstm" | "rnn"
type LSTMState = { h: Types.Tensor, c: Types.Tensor }

local function makeCell(cellType: CellType, inputDim: number, hiddenDim: number)
	if cellType == "lstm" then
		return LSTM(inputDim, hiddenDim)
	elseif cellType == "rnn" then
		return RNN(inputDim, hiddenDim)
	end
	return GRU(inputDim, hiddenDim)
end

local function sliceTimeStep(x: Types.Tensor, t: number, inputDim: number, batch: number): Types.Tensor
	local step = Tensor.slice(x, 2, t, t)
	return Tensor.reshape(step, { inputDim, batch })
end

local function initStates(self, batch: number, dtype: Types.DType, requiresGrad: boolean)
	local states = table.create(#self._cells)
	for i = 1, #self._cells do
		if self._cellType == "lstm" then
			states[i] = {
				h = Tensor.zeros({ self._hiddenDim, batch }, dtype, requiresGrad),
				c = Tensor.zeros({ self._hiddenDim, batch }, dtype, requiresGrad),
			}
		else
			states[i] = Tensor.zeros({ self._hiddenDim, batch }, dtype, requiresGrad)
		end
	end
	return states
end

local function forward(self, x: Types.Tensor): Types.Tensor
	assert(#x._shape == 3, "SequenceClassifier: expected input shape {features, time, batch}")
	local inputDim, timeSteps, batch = x._shape[1], x._shape[2], x._shape[3]
	assert(inputDim == self._inputDim, "SequenceClassifier: inputDim mismatch")
	local states = initStates(self, batch, x._dtype, x._requiresGrad)
	local last
	for t = 1, timeSteps do
		local step = sliceTimeStep(x, t, inputDim, batch)
		local layerInput: Types.Tensor = step
		for layerIdx, cell in ipairs(self._cells) do
			if self._cellType == "lstm" then
				local state = states[layerIdx] :: LSTMState
				local nextH, nextC = cell:forward(layerInput, state.h, state.c)
				states[layerIdx] = { h = nextH, c = nextC }
				layerInput = nextH
			else
				local prev = states[layerIdx] :: Types.Tensor
				local nextH = cell:forward(layerInput, prev)
				states[layerIdx] = nextH
				layerInput = nextH
			end
			local drop = self._dropoutBetween[layerIdx]
			if drop and layerIdx < #self._cells then
				layerInput = drop:forward(layerInput)
			end
		end
		last = layerInput
	end
	assert(last, "SequenceClassifier: sequence length must be >= 1")
	local out = last
	if self._norm then
		out = self._norm:forward(out)
	end
	if self._finalDropout then
		out = self._finalDropout:forward(out)
	end
	out = self._head:forward(out)
	if self._finalAct then
		out = self._finalAct:forward(out)
	end
	return out
end

local function parameters(self): { Types.Tensor }
	local params = {}
	for _, cell in ipairs(self._cells) do
		Utils.collectParams(params, cell)
	end
	if self._norm then
		Utils.collectParams(params, self._norm)
	end
	Utils.collectParams(params, self._head)
	return params
end

local function train(self, mode: boolean?)
	Utils.trainAll(self._trainables, mode)
end

local function SequenceClassifier(cfg: SequenceClassifierConfig): Types.Module<Types.Tensor, Types.Tensor>
	assert(cfg.inputDim and cfg.hiddenDim and cfg.numClasses, "SequenceClassifier: missing config fields")
	local numLayers = math.max(1, cfg.numLayers or 1)
	local cellType = (cfg.cell or "gru") :: CellType

	local layers = table.create(numLayers)
	local inDim = cfg.inputDim
	for i = 1, numLayers do
		layers[i] = makeCell(cellType, inDim, cfg.hiddenDim)
		inDim = cfg.hiddenDim
	end

	local dropoutBetween = {}
	local dropProb = math.clamp(cfg.dropout or 0, 0, 0.95)
	if dropProb > 0 and numLayers > 1 then
		for i = 1, numLayers - 1 do
			dropoutBetween[i] = Dropout(dropProb)
		end
	end
	local finalDropout = if dropProb > 0 then Dropout(dropProb) else nil
	local norm = if cfg.layerNorm then LayerNorm(cfg.hiddenDim) else nil
	local head = Linear(cfg.hiddenDim, cfg.numClasses)
	local finalAct = if cfg.finalActivation then Utils.makeActivationModule(cfg.finalActivation) else nil

	local trainables: { any } = {}
	for _, cell in ipairs(layers) do
		table.insert(trainables, cell)
	end
	for _, drop in ipairs(dropoutBetween) do
		table.insert(trainables, drop)
	end
	if finalDropout then
		table.insert(trainables, finalDropout)
	end
	if norm then
		table.insert(trainables, norm)
	end
	table.insert(trainables, head)
	if finalAct then
		table.insert(trainables, finalAct)
	end

	return {
		_inputDim = cfg.inputDim,
		_hiddenDim = cfg.hiddenDim,
		_cells = layers,
		_cellType = cellType,
		_dropoutBetween = dropoutBetween,
		_finalDropout = finalDropout,
		_norm = norm,
		_head = head,
		_finalAct = finalAct,
		_trainables = trainables,
		forward = forward,
		parameters = parameters,
		train = train,
	}
end

return {
	new = SequenceClassifier,
}


--!native
--!optimize 2
local Types = require("../Types")
local Utils = require("./Utils")
local Linear = require("../nn/Linear")
local LayerNorm = require("../nn/LayerNorm")
local Dropout = require("../nn/Dropout")
local Sequential = require("../nn/Sequential")
local Attention = require("../nn/Attention")

export type TransformerEncoderConfig = {
	embedDim: number,
	numHeads: number,
	mlpDim: number?,
	numLayers: number?,
	inputDim: number?,
	dropout: number?,
	attentionDropout: number?,
	activation: Utils.ActivationName?,
	finalLayerNorm: boolean?,
	weightInit: ((number, number) -> number)?,
}

type Block = {
	norm1: Types.Module<Types.Tensor, Types.Tensor>,
	attn: Types.Module<Types.Tensor, Types.Tensor>,
	drop1: Types.Module<Types.Tensor, Types.Tensor>?,
	norm2: Types.Module<Types.Tensor, Types.Tensor>,
	ff: Types.Module<Types.Tensor, Types.Tensor>,
	trainables: { any },
}

local function buildFeedForward(
	embedDim: number,
	mlpDim: number,
	activation: Utils.ActivationName?,
	dropProb: number,
	init: ((number, number) -> number)?
): (Types.Module<Types.Tensor, Types.Tensor>, { any })
	local layers = {}
	local trainables: { any } = {}
	local function insert(m)
		table.insert(layers, m)
		table.insert(trainables, m)
	end
	insert(Linear(embedDim, mlpDim, init))
	insert(Utils.makeActivationModule(activation or "gelu"))
	if dropProb > 0 then
		insert(Dropout(dropProb))
	end
	insert(Linear(mlpDim, embedDim, init))
	if dropProb > 0 then
		insert(Dropout(dropProb))
	end
	return Sequential(layers), trainables
end

local function buildBlock(embedDim: number, cfg: TransformerEncoderConfig): Block
	local drop = math.clamp(cfg.dropout or 0, 0, 0.95)
	local norm1 = LayerNorm(embedDim)
	local attn = Attention(embedDim, cfg.numHeads, {
		attentionDropout = cfg.attentionDropout or drop,
	})
	local drop1 = if drop > 0 then Dropout(drop) else nil
	local norm2 = LayerNorm(embedDim)
	local mlpDim = cfg.mlpDim or (embedDim * 4)
	local ff, ffTrainables = buildFeedForward(embedDim, mlpDim, cfg.activation, drop, cfg.weightInit)
	local trainables: { any } = { norm1, attn, norm2 }
	if drop1 then
		table.insert(trainables, drop1)
	end
	for _, m in ipairs(ffTrainables) do
		table.insert(trainables, m)
	end

	return {
		norm1 = norm1,
		attn = attn,
		drop1 = drop1,
		norm2 = norm2,
		ff = ff,
		trainables = trainables,
	}
end

local function buildEncoder(cfg: TransformerEncoderConfig): Types.Module<Types.Tensor, Types.Tensor>
	assert(cfg.embedDim and cfg.numHeads, "TransformerEncoder: embedDim/numHeads required")
	local numLayers = math.max(1, cfg.numLayers or 4)
	local blocks = table.create(numLayers)
	local trainables: { any } = {}
	for i = 1, numLayers do
		local block = buildBlock(cfg.embedDim, cfg)
		blocks[i] = block
		for _, m in ipairs(block.trainables) do
			table.insert(trainables, m)
		end
	end

	local inputProj = if cfg.inputDim and cfg.inputDim ~= cfg.embedDim
		then Linear(cfg.inputDim, cfg.embedDim, cfg.weightInit)
		else nil
	local finalNorm = if cfg.finalLayerNorm ~= false then LayerNorm(cfg.embedDim) else nil
	if inputProj then
		table.insert(trainables, inputProj)
	end
	if finalNorm then
		table.insert(trainables, finalNorm)
	end

	local function forward(_, x: Types.Tensor): Types.Tensor
		local h = if inputProj then inputProj:forward(x) else x
		for _, block in ipairs(blocks) do
			local y = block.norm1:forward(h)
			y = block.attn:forward(y)
			if block.drop1 then
				y = block.drop1:forward(y)
			end
			h = Utils.addTensors(h, y)

			local z = block.norm2:forward(h)
			z = block.ff:forward(z)
			h = Utils.addTensors(h, z)
		end
		if finalNorm then
			h = finalNorm:forward(h)
		end
		return h
	end

	local function parameters(_): { Types.Tensor }
		local params = {}
		if inputProj then
			Utils.collectParams(params, inputProj)
		end
		for _, block in ipairs(blocks) do
			Utils.collectParams(params, block.norm1)
			Utils.collectParams(params, block.attn)
			if block.drop1 then
				Utils.collectParams(params, block.drop1)
			end
			Utils.collectParams(params, block.norm2)
			Utils.collectParams(params, block.ff)
		end
		if finalNorm then
			Utils.collectParams(params, finalNorm)
		end
		return params
	end

	local function train(_, mode: boolean?)
		Utils.trainAll(trainables, mode)
	end

	return {
		forward = forward,
		parameters = parameters,
		train = train,
	}
end

return {
	new = buildEncoder,
}


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")
local Activations = require("../nn/Activations")
local Util = require("../Util")

local assert = Util.Assert

export type ActivationName = "relu" | "gelu" | "tanh" | "silu" | "sigmoid" | "mish" | "elu" | "leaky_relu" | "selu"

type Module = Types.Module<Types.Tensor, Types.Tensor>

local Utils = {}

local ActivationLookup: { [string]: (Types.Tensor) -> Types.Tensor } = {
	relu = Activations.ReLU,
	gelu = Activations.GELU,
	tanh = Activations.Tanh,
	silu = Activations.SiLU,
	sigmoid = Activations.Sigmoid,
	mish = Activations.Mish,
	elu = function(x: Types.Tensor): Types.Tensor
		return Activations.ELU(x)
	end,
	leaky_relu = function(x: Types.Tensor): Types.Tensor
		return Activations.LeakyReLU(x, 0.01)
	end,
	selu = Activations.SeLU,
}

local function normalizeKey(name: string?): string
	if not name then
		return "relu"
	end
	return string.lower(name)
end

function Utils.getActivationFn(name: ActivationName?): (Types.Tensor) -> Types.Tensor
	local fn = ActivationLookup[normalizeKey(name)]
	assert(fn, `Unsupported activation '{name}'.`)
	return fn
end

function Utils.makeActivationModule(name: ActivationName?): Module
	local fn = Utils.getActivationFn(name)
	return {
		forward = function(_, x: Types.Tensor): Types.Tensor
			return fn(x)
		end,
		parameters = function(_): { Types.Tensor }
			return {}
		end,
		train = function(_) end,
	}
end

function Utils.collectParams(dst: { Types.Tensor }, module: Module?)
	if not module then
		return
	end
	local ok, params = pcall(function()
		return module:parameters()
	end)
	if ok and params then
		for _, p in ipairs(params) do
			table.insert(dst, p)
		end
	end
end

function Utils.collectFromList(dst: { Types.Tensor }, modules: { Module? })
	for _, m in ipairs(modules) do
		Utils.collectParams(dst, m)
	end
end

function Utils.trainAll(modules: { any }, mode: boolean?)
	for _, module in ipairs(modules) do
		if type(module) == "table" then
			local trainFn = (module :: any).train
			if type(trainFn) == "function" then
				trainFn(module, mode)
			end
		end
	end
end

function Utils.addTensors(a: Types.Tensor, b: Types.Tensor): Types.Tensor
	assert(#a._shape == #b._shape, "addTensors: shape rank mismatch")
	for i = 1, #a._shape do
		assert(a._shape[i] == b._shape[i], "addTensors: shape mismatch")
	end
	local requiresGrad = (a._requiresGrad or b._requiresGrad) == true
	local out = Tensor.zeros(a._shape, a._dtype, requiresGrad)
	task.desynchronize()
	for i = 1, #a._storage do
		out._storage[i] = a._storage[i] + b._storage[i]
	end
	task.synchronize()
	if not requiresGrad then
		return out
	end
	return Tape.makeNode(out, { a, b }, function(gy: Types.Tensor)
		task.desynchronize()
		if a._requiresGrad then
			a._grad = a._grad or Tensor.zeros(a._shape)
			for i = 1, #gy._storage do
				(a._grad :: Types.Tensor)._storage[i] += gy._storage[i]
			end
		end
		if b._requiresGrad then
			b._grad = b._grad or Tensor.zeros(b._shape)
			for i = 1, #gy._storage do
				(b._grad :: Types.Tensor)._storage[i] += gy._storage[i]
			end
		end
		task.synchronize()
	end)
end

return Utils


--!native
--!optimize 2
local Tape = require("../autograd/Tape")
local Tensor = require("../Tensor")
local Types = require("../Types")
local Util = require("../Util")

local assert = Util.Assert

local A = {}

local function prep(x: Types.Tensor): Types.Tensor
	return Tensor.zeros(x._shape, x._dtype, x._requiresGrad)
end

local function ensureGrad(x: Types.Tensor): Types.Tensor?
	if not x._grad then
		x._grad = Tensor.zeros(x._shape)
	end
	return x._grad
end

function A.ReLU(x: Types.Tensor): Types.Tensor
	return Tape.relu(x)
end

function A.Sigmoid(x: Types.Tensor): Types.Tensor
	task.desynchronize()
	local y = prep(x)
	local xs, ys = x._storage, y._storage
	for i = 1, #xs do
		ys[i] = 1 / (1 + math.exp(-xs[i]))
	end
	task.synchronize()
	return Tape.makeNode(y, { x }, function(gy)
		if not x._requiresGrad then
			return
		end
		local gx = ensureGrad(x)
		task.desynchronize()
		local gxs, gys = gx._storage, gy._storage
		for i = 1, #ys do
			local s = ys[i]
			gxs[i] += gys[i] * s * (1 - s)
		end
		task.synchronize()
	end)
end

function A.Tanh(x: Types.Tensor): Types.Tensor
	task.desynchronize()
	local y = prep(x)
	local xs, ys = x._storage, y._storage
	for i = 1, #xs do
		ys[i] = math.tanh(xs[i])
	end
	task.synchronize()
	return Tape.makeNode(y, { x }, function(gy)
		if not x._requiresGrad then
			return
		end
		local gx = ensureGrad(x)
		task.desynchronize()
		local gxs, gys = gx._storage, gy._storage
		for i = 1, #ys do
			local t = ys[i]
			gxs[i] += gys[i] * (1 - t * t)
		end
		task.synchronize()
	end)
end

function A.GELU(x: Types.Tensor): Types.Tensor
	task.desynchronize()
	local y = prep(x)
	local xs, ys = x._storage, y._storage
	local c = math.sqrt(2 / math.pi)
	local k = 0.044715

	for i = 1, #xs do
		local v = xs[i]
		local cube = v * v * v
		local inner = c * (v + k * cube)
		ys[i] = 0.5 * v * (1 + math.tanh(inner))
	end
	task.synchronize()
	return Tape.makeNode(y, { x }, function(gy)
		if not x._requiresGrad then
			return
		end
		local gx = ensureGrad(x)
		task.desynchronize()
		local gxs, gys = gx._storage, gy._storage
		local c2 = math.sqrt(2 / math.pi)

		for i = 1, #xs do
			local v = xs[i]
			local v2 = v * v
			local inner = c2 * (v + 0.044715 * v * v2)
			local t = math.tanh(inner)
			local sech2 = 1 - t * t
			local d_inner = c2 * (1 + 3 * 0.044715 * v2)

			local d = 0.5 * (1 + t) + 0.5 * v * sech2 * d_inner
			gxs[i] += gys[i] * d
		end
		task.synchronize()
	end)
end

function A.LeakyReLU(x: Types.Tensor, alpha: number?): Types.Tensor
	local a = alpha or 0.01
	task.desynchronize()
	local y = prep(x)
	local xs, ys = x._storage, y._storage
	for i = 1, #xs do
		local v = xs[i]
		ys[i] = (v > 0) and v or (a * v)
	end
	task.synchronize()
	return Tape.makeNode(y, { x }, function(gy)
		if not x._requiresGrad then
			return
		end
		local gx = ensureGrad(x)
		task.desynchronize()
		local gxs, gys = gx._storage, gy._storage
		for i = 1, #xs do
			local v = xs[i]
			local slope = (v > 0) and 1 or a
			gxs[i] += gys[i] * slope
		end
		task.synchronize()
	end)
end

function A.Softplus(x: Types.Tensor): Types.Tensor
	task.desynchronize()
	local y = prep(x)
	local xs, ys = x._storage, y._storage
	for i = 1, #xs do
		ys[i] = math.log(1 + math.exp(xs[i]))
	end
	task.synchronize()
	return Tape.makeNode(y, { x }, function(gy)
		if not x._requiresGrad then
			return
		end
		local gx = ensureGrad(x)
		task.desynchronize()
		local gxs, gys = gx._storage, gy._storage
		for i = 1, #xs do
			local v = xs[i]
			local s = 1 / (1 + math.exp(-v)) -- sigmoid
			gxs[i] += gys[i] * s
		end
		task.synchronize()
	end)
end

function A.ELU(x: Types.Tensor, alpha: number?): Types.Tensor
	local a = alpha or 1.0
	task.desynchronize()
	local y = prep(x)
	local xs, ys = x._storage, y._storage
	for i = 1, #xs do
		local v = xs[i]
		ys[i] = (v > 0) and v or (a * (math.exp(v) - 1))
	end
	task.synchronize()
	return Tape.makeNode(y, { x }, function(gy)
		if not x._requiresGrad then
			return
		end
		local gx = ensureGrad(x)
		task.desynchronize()
		local gxs, gys = gx._storage, gy._storage
		for i = 1, #xs do
			local v = xs[i]
			local d = (v > 0) and 1 or (ys[i] + a)
			gxs[i] += gys[i] * d
		end
		task.synchronize()
	end)
end

function A.SiLU(x: Types.Tensor): Types.Tensor
	task.desynchronize()
	local y = prep(x)
	local xs, ys = x._storage, y._storage
	for i = 1, #xs do
		local v = xs[i]
		local s = 1 / (1 + math.exp(-v))
		ys[i] = v * s
	end
	task.synchronize()
	return Tape.makeNode(y, { x }, function(gy)
		if not x._requiresGrad then
			return
		end
		local gx = ensureGrad(x)
		task.desynchronize()
		local gxs, gys = gx._storage, gy._storage
		for i = 1, #xs do
			local v = xs[i]
			local s = 1 / (1 + math.exp(-v))
			local d = s + v * s * (1 - s)
			gxs[i] += gys[i] * d
		end
		task.synchronize()
	end)
end

function A.Mish(x: Types.Tensor): Types.Tensor
	task.desynchronize()
	local y = prep(x)
	local xs, ys = x._storage, y._storage
	local t = 20 -- threshold
	for i = 1, #xs do
		local v = xs[i]
		if v > t then
			-- softplus(x)
			ys[i] = v * math.tanh(v)
		else
			local sp = math.log(1 + math.exp(v))
			ys[i] = v * math.tanh(sp)
		end
	end
	task.synchronize()
	return Tape.makeNode(y, { x }, function(gy)
		if not x._requiresGrad then
			return
		end
		local gx = ensureGrad(x)
		task.desynchronize()
		local gxs, gys = gx._storage, gy._storage
		for i = 1, #xs do
			local v = xs[i]
			local d
			if v > t then
				local tanh_v = math.tanh(v)
				d = tanh_v + v * (1 - tanh_v * tanh_v)
			else
				local exp_v = math.exp(v)
				local sp = math.log(1 + exp_v)
				local tanh_sp = math.tanh(sp)
				local sigmoid = 1 / (1 + math.exp(-v))
				d = tanh_sp + v * (1 - tanh_sp * tanh_sp) * sigmoid
			end
			gxs[i] += gys[i] * d
		end
		task.synchronize()
	end)
end

function A.SeLU(x: Types.Tensor, alpha: number?, lambda: number?): Types.Tensor
	local a = alpha or 1.6732632423543772
	local l = lambda or 1.0507009873554805
	task.desynchronize()
	local y = prep(x)
	local xs, ys = x._storage, y._storage
	for i = 1, #xs do
		local v = xs[i]
		ys[i] = (v > 0) and (l * v) or (l * a * (math.exp(v) - 1))
	end
	task.synchronize()
	return Tape.makeNode(y, { x }, function(gy)
		if not x._requiresGrad then
			return
		end
		local gx = ensureGrad(x)
		task.desynchronize()
		local gxs, gys = gx._storage, gy._storage
		local la = l * a
		for i = 1, #xs do
			local v = xs[i]
			local d = (v > 0) and l or (la * math.exp(v))
			gxs[i] += gys[i] * d
		end
		task.synchronize()
	end)
end

function A.SwiGLU(a: Types.Tensor, b: Types.Tensor): Types.Tensor
	assert(#a._storage == #b._storage, "SwiGLU: size mismatch")
	task.desynchronize()
	local y = Tensor.zeros(a._shape, a._dtype, a._requiresGrad or b._requiresGrad)
	local as, bs, ys = a._storage, b._storage, y._storage
	for i = 1, #as do
		local val_a = as[i]
		local s = 1 / (1 + math.exp(-val_a))
		ys[i] = (val_a * s) * bs[i]
	end
	task.synchronize()
	return Tape.makeNode(y, { a, b }, function(gy)
		if not (a._requiresGrad or b._requiresGrad) then
			return
		end

		local ga = a._requiresGrad and ensureGrad(a)
		local gb = b._requiresGrad and ensureGrad(b)

		task.desynchronize()
		local gas = ga and ga._storage
		local gbs = gb and gb._storage
		local gys = gy._storage

		for i = 1, #as do
			local val_a = as[i]
			local val_b = bs[i]
			local g = gys[i]

			local s = 1 / (1 + math.exp(-val_a))
			local swish = val_a * s

			-- d/da (swish * b) = b * (swish + x * s * (1-s)) -> actually d/dx (x*s) = s + x*s*(1-s)
			local d_swish = s + val_a * s * (1 - s)

			if gas then
				gas[i] += g * (val_b * d_swish)
			end
			if gbs then
				gbs[i] += g * swish
			end
		end
		task.synchronize()
	end)
end

function A.SwiGLUSplit(x: Types.Tensor, hidden: number?): Types.Tensor
	local rows, B = x._shape[1], x._shape[2]
	local H = math.floor(rows / 2)
	assert(rows == 2 * H, "SwiGLUSplit: first dim must be even")
	if hidden then
		assert(hidden == H, "SwiGLUSplit: hidden mismatch")
	end

	task.desynchronize()
	local y = Tensor.zeros({ H, B }, x._dtype, x._requiresGrad)
	local xs, ys = x._storage, y._storage
	for j = 1, B do
		-- a is at x[(0...H-1)*B + j]
		-- b is at x[(H...2H-1)*B + j]
		for i = 1, H do
			local idx_a = (i - 1) * B + j
			local idx_b = (i + H - 1) * B + j

			local val_a = xs[idx_a]
			local val_b = xs[idx_b]

			local s = 1 / (1 + math.exp(-val_a))
			ys[(i - 1) * B + j] = (val_a * s) * val_b
		end
	end
	task.synchronize()
	return Tape.makeNode(y, { x }, function(gy)
		if not x._requiresGrad then
			return
		end
		local gx = ensureGrad(x)
		task.desynchronize()
		local gxs = gx._storage
		local gys = gy._storage
		for j = 1, B do
			for i = 1, H do
				local idx_a = (i - 1) * B + j
				local idx_b = (i + H - 1) * B + j
				local idx_y = (i - 1) * B + j

				local val_a = xs[idx_a]
				local val_b = xs[idx_b]
				local g = gys[idx_y]

				local s = 1 / (1 + math.exp(-val_a))
				local swish = val_a * s
				local d_swish = s + val_a * s * (1 - s)

				gxs[idx_a] += g * (val_b * d_swish)
				gxs[idx_b] += g * swish
			end
		end
		task.synchronize()
	end)
end

return A


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")
local Linear = require("./Linear")
local SoftmaxOps = require("../ops/Softmax")
local Util = require("../Util")

local assert = Util.Assert

type TensorT = Types.Tensor
type Module = Types.Module<TensorT, TensorT>
type Random = typeof(Random.new())

export type AttentionOptions = {
	qkvInit: ((number, number) -> number)?,
	projInit: ((number, number) -> number)?,
	attentionDropout: number?, -- dropout prob applied to attention weights during training
	dropoutRNG: Random?, -- optional RNG source for dropout sampling
	causal: boolean?, -- enable built-in causal masking
}
local NEG_INF = -1e9

local Attention = {}
Attention.__index = Attention

local function scaleTensor(x: TensorT, factor: number): TensorT
	local y = Tensor.zeros(x._shape, x._dtype, x._requiresGrad)
	task.desynchronize()
	for i = 1, #x._storage do
		y._storage[i] = x._storage[i] * factor
	end
	task.synchronize()
	if not x._requiresGrad then
		return y
	end
	return Tape.makeNode(y, { x }, function(gy: TensorT)
		if not x._requiresGrad then
			return
		end
		x._grad = x._grad or Tensor.zeros(x._shape)
		local grad = x._grad :: TensorT
		task.desynchronize()
		for i = 1, #gy._storage do
			grad._storage[i] += gy._storage[i] * factor
		end
		task.synchronize()
	end)
end

local function addTensors(a: TensorT, b: TensorT): TensorT
	assert(#a._shape == #b._shape, "addTensors: shape rank mismatch")
	for i = 1, #a._shape do
		assert(a._shape[i] == b._shape[i], "addTensors: shape mismatch")
	end
	local requiresGrad = (a._requiresGrad or b._requiresGrad) == true
	local out = Tensor.zeros(a._shape, a._dtype, requiresGrad)
	task.desynchronize()
	for i = 1, #a._storage do
		out._storage[i] = a._storage[i] + b._storage[i]
	end
	task.synchronize()
	if not requiresGrad then
		return out
	end
	return Tape.makeNode(out, { a, b }, function(gy: TensorT)
		task.desynchronize()
		if a._requiresGrad then
			a._grad = a._grad or Tensor.zeros(a._shape)
			local gradA = a._grad :: TensorT
			for i = 1, #gy._storage do
				gradA._storage[i] += gy._storage[i]
			end
		end
		if b._requiresGrad then
			b._grad = b._grad or Tensor.zeros(b._shape)
			local gradB = b._grad :: TensorT
			for i = 1, #gy._storage do
				gradB._storage[i] += gy._storage[i]
			end
		end
		task.synchronize()
	end)
end

local function dropoutTensor(x: TensorT, dropProb: number, rng: Random?): TensorT
	if dropProb <= 0 then
		return x
	end
	local keep = 1 - dropProb
	assert(keep > 0, "dropoutTensor: keep probability must be > 0")
	local mask = table.create(#x._storage)
	local y = Tensor.zeros(x._shape, x._dtype, x._requiresGrad)
	task.desynchronize()
	for i = 1, #x._storage do
		local keepSample = ((rng or Random.new()):NextNumber(0, 1) < keep) and 1 or 0
		mask[i] = keepSample
		y._storage[i] = x._storage[i] * keepSample / keep
	end
	task.synchronize()
	if not x._requiresGrad then
		return y
	end
	return Tape.makeNode(y, { x }, function(gy: TensorT)
		if not x._requiresGrad then
			return
		end
		x._grad = x._grad or Tensor.zeros(x._shape)
		local grad = x._grad :: TensorT
		task.desynchronize()
		for i = 1, #gy._storage do
			grad._storage[i] += gy._storage[i] * (mask[i] / keep)
		end
		task.synchronize()
	end)
end

local function softmaxRows(x: TensorT): TensorT
	local rows, cols = x._shape[1], x._shape[2]
	local y = SoftmaxOps.forward(x)
	if not x._requiresGrad then
		return y
	end
	return Tape.makeNode(y, { x }, function(gy: TensorT)
		if not x._requiresGrad then
			return
		end
		x._grad = x._grad or Tensor.zeros(x._shape)
		local grad = x._grad :: TensorT
		task.desynchronize()
		local yStore = y._storage
		local gyStore = gy._storage
		local gxStore = grad._storage
		for i = 1, rows do
			local offset = (i - 1) * cols
			local dot = 0
			for j = 1, cols do
				local idx = offset + j
				dot += yStore[idx] * gyStore[idx]
			end
			for j = 1, cols do
				local idx = offset + j
				gxStore[idx] += yStore[idx] * (gyStore[idx] - dot)
			end
		end
		task.synchronize()
	end)
end

local function concatHeads(heads: { TensorT }, headDim: number, tokens: number): TensorT
	assert(#heads > 0, "concatHeads: expected at least one head")
	local embedDim = headDim * #heads
	local requiresGrad = false
	for _, h in heads do
		if h._requiresGrad then
			requiresGrad = true
			break
		end
	end
	local y = Tensor.zeros({ embedDim, tokens }, heads[1]._dtype, requiresGrad)
	task.desynchronize()
	local yStore = y._storage
	for idxHead = 1, #heads do
		local h = heads[idxHead]
		assert(h._shape[1] == headDim, "concatHeads: head dim mismatch")
		assert(h._shape[2] == tokens, "concatHeads: token dim mismatch")
		local hStore = h._storage
		local rowOffset = (idxHead - 1) * headDim
		for i = 1, headDim do
			for j = 1, tokens do
				local dstRow = rowOffset + i
				local dstIdx = (dstRow - 1) * tokens + j
				local srcIdx = (i - 1) * tokens + j
				yStore[dstIdx] = hStore[srcIdx]
			end
		end
	end
	task.synchronize()
	if not requiresGrad then
		return y
	end
	return Tape.makeNode(y, heads, function(gy: TensorT)
		task.desynchronize()
		for idxHead = 1, #heads do
			local h = heads[idxHead]
			if h._requiresGrad then
				h._grad = h._grad or Tensor.zeros(h._shape)
				local hg = (h._grad :: TensorT)._storage
				local rowOffset = (idxHead - 1) * headDim
				for i = 1, headDim do
					for j = 1, tokens do
						local dstRow = rowOffset + i
						local gyIdx = (dstRow - 1) * tokens + j
						local srcIdx = (i - 1) * tokens + j
						hg[srcIdx] += gy._storage[gyIdx]
					end
				end
			end
		end
		task.synchronize()
	end)
end

type DropoutConfig = {
	prob: number,
	rng: Random?,
	enabled: boolean?,
}

local function scaledDotProduct(
	q: TensorT,
	k: TensorT,
	v: TensorT,
	mask: TensorT?,
	scaleOverride: number?,
	dropCfg: DropoutConfig?
): (TensorT, TensorT)
	local headDim = q._shape[1]
	local queryLen = q._shape[2]
	local keyLen = k._shape[2]
	assert(k._shape[1] == headDim, "scaledDotProduct: Q/K dim mismatch")
	assert(v._shape[1] == headDim, "scaledDotProduct: Q/V dim mismatch")
	assert(v._shape[2] == keyLen, "scaledDotProduct: K/V length mismatch")
	local scale = scaleOverride or (1 / math.sqrt(headDim))
	local scores = Tape.matmul(Tensor.transpose(q), k) -- {queryLen, keyLen}
	scores = scaleTensor(scores, scale)
	if mask then
		assert(mask._shape[1] == queryLen and mask._shape[2] == keyLen, "scaledDotProduct: mask shape mismatch")
		scores = addTensors(scores, mask)
	end
	local attn = softmaxRows(scores)
	if dropCfg and dropCfg.prob > 0 and dropCfg.enabled ~= false then
		attn = dropoutTensor(attn, dropCfg.prob, dropCfg.rng)
	end
	local context = Tensor.transpose(Tape.matmul(attn, Tensor.transpose(v)))
	return context, attn
end

function Attention.new(embedDim: number, numHeads: number, opts: AttentionOptions?): Module
	assert(numHeads >= 1, "Attention: numHeads must be >= 1")
	assert(embedDim % numHeads == 0, "Attention: embedDim must be divisible by numHeads")
	opts = opts or {}
	local options = opts :: AttentionOptions
	local headDim = math.floor(embedDim / numHeads)
	local self = setmetatable({
		_embedDim = embedDim,
		_numHeads = numHeads,
		_headDim = headDim,
		_scale = 1 / math.sqrt(headDim),
		_qkv = Linear(embedDim, embedDim * 3, options.qkvInit),
		_proj = Linear(embedDim, embedDim, options.projInit),
		_attentionDropout = math.clamp(options.attentionDropout or 0, 0, 0.999),
		_dropoutRNG = options.dropoutRNG or Random.new(),
		_train = true,
		_causal = options.causal == true,
		_causalMaskCache = {} :: { [Types.DType]: { [number]: TensorT } },
		_lastAttn = nil,
	}, Attention)
	return self
end

function Attention:forward(x: TensorT, mask: TensorT?): TensorT
	local tokens = x._shape[2]
	local embedDim = self._embedDim
	local qkv = self._qkv:forward(x) -- {3*D, tokens}
	assert(qkv._shape[1] == embedDim * 3, "Attention: unexpected projection shape")
	local q = Tensor.narrow(qkv, 1, 1, embedDim)
	local k = Tensor.narrow(qkv, 1, embedDim + 1, embedDim)
	local v = Tensor.narrow(qkv, 1, 2 * embedDim + 1, embedDim)
	local effMask = mask
	if not effMask and self._causal then
		effMask = self:_getCausalMask(tokens, q._dtype)
	end
	local dropCfg: DropoutConfig?
	if self._attentionDropout > 0 and self._train then
		dropCfg = {
			prob = self._attentionDropout,
			rng = self._dropoutRNG,
			enabled = true,
		}
	end
	local headOutputs = table.create(self._numHeads)
	local attnWeights = table.create(self._numHeads)
	for h = 0, self._numHeads - 1 do
		local offset = h * self._headDim + 1
		local qh = Tensor.narrow(q, 1, offset, self._headDim)
		local kh = Tensor.narrow(k, 1, offset, self._headDim)
		local vh = Tensor.narrow(v, 1, offset, self._headDim)
		local ctx, weights = scaledDotProduct(qh, kh, vh, effMask, self._scale, dropCfg)
		headOutputs[h + 1] = ctx
		attnWeights[h + 1] = weights
	end
	self._lastAttn = attnWeights
	local merged = concatHeads(headOutputs, self._headDim, tokens)
	return self._proj:forward(merged)
end

function Attention:parameters(): { TensorT }
	local params = {}
	for _, t in self._qkv:parameters() do
		table.insert(params, t)
	end
	for _, t in self._proj:parameters() do
		table.insert(params, t)
	end
	return params
end

function Attention:getLastAttention(): { TensorT }?
	return self._lastAttn
end

function Attention:train(mode: boolean?)
	self._train = (mode ~= false)
end

function Attention:setCausal(enabled: boolean?)
	self._causal = enabled ~= false
end

function Attention:_getCausalMask(tokens: number, dtype: Types.DType?): TensorT
	dtype = dtype or "f64"
	local cacheByDtype = self._causalMaskCache[dtype]
	if not cacheByDtype then
		cacheByDtype = {}
		self._causalMaskCache[dtype] = cacheByDtype
	end
	local mask = cacheByDtype[tokens]
	if mask then
		return mask
	end
	local newMask = Tensor.zeros({ tokens, tokens }, dtype, false)
	task.desynchronize()
	local store = newMask._storage
	for i = 1, tokens do
		local offset = (i - 1) * tokens
		for j = i + 1, tokens do
			store[offset + j] = NEG_INF
		end
	end
	task.synchronize()
	cacheByDtype[tokens] = newMask
	return newMask
end

local exported = {
	new = function(embedDim: number, numHeads: number, opts: AttentionOptions?): Module
		return Attention.new(embedDim, numHeads, opts)
	end,
	scaledDotProduct = scaledDotProduct,
}

setmetatable(exported, {
	__call = function(_, embedDim: number, numHeads: number, opts: AttentionOptions?): Module
		return Attention.new(embedDim, numHeads, opts)
	end,
})

return exported


--!native
--!optimize 2
local Avg = require("../ops/AvgPool2d")
local Types = require("../Types")
return function(KH: number, KW: number, stride: number)
	return {
		forward = function(_, x: Types.Tensor): Types.Tensor
			return Avg(x, KH, KW, stride)
		end,
		parameters = function(_)
			return {}
		end,
		train = function(_) end,
	}
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")
local Util = require("../Util")
local assert = Util.Assert

-- BatchNorm1d (shape {D,B})
--@param D number
--@param eps number?
--@param momentum number?
--@return table
return function(D: number, eps: number?, momentum: number?): Types.Module<Types.Tensor, Types.Tensor>
	eps = eps or 1e-5
	momentum = momentum or 0.1
	local gamma = Tensor.fromArray(table.create(D, 1), { D, 1 }, "f64", true)
	local beta = Tensor.zeros({ D, 1 }, "f64", true)
	local runningMean = Tensor.zeros({ D, 1 })
	local runningVar = Tensor.fromArray(table.create(D, 1), { D, 1 })
	return {
		_train = true,
		train = function(self, mode: boolean)
			self._train = mode == nil and true or mode
		end,
		forward = function(self, x: Types.Tensor): Types.Tensor
			local D_, B = x._shape[1], x._shape[2]
			assert(D_ == D, "BN1d D mismatch")
			local xc = Tensor.zeros(x._shape, x._dtype, x._requiresGrad)
			local meanUsed: Types.Tensor
			local varUsed: Types.Tensor
			if self._train then
				meanUsed = Tensor.zeros({ D, 1 })
				varUsed = Tensor.zeros({ D, 1 })
				task.desynchronize()
				for i = 1, D do
					local sum = 0
					local base = (i - 1) * B
					for j = 1, B do
						sum += x._storage[base + j]
					end
					local mean = sum / B
					meanUsed._storage[i] = mean
					local vsum = 0
					for j = 1, B do
						local idx = base + j
						local centered = x._storage[idx] - mean
						xc._storage[idx] = centered
						vsum += centered * centered
					end
					varUsed._storage[i] = vsum / B
				end
				task.synchronize()
				task.desynchronize()
				for i = 1, D do
					runningMean._storage[i] = (1 - (momentum :: number)) * runningMean._storage[i]
						+ (momentum :: number) * meanUsed._storage[i]
					runningVar._storage[i] = (1 - (momentum :: number)) * runningVar._storage[i]
						+ (momentum :: number) * varUsed._storage[i]
				end
				task.synchronize()
			else
				meanUsed = runningMean
				varUsed = runningVar
				task.desynchronize()
				for i = 1, D do
					local base = (i - 1) * B
					local mean = meanUsed._storage[i]
					for j = 1, B do
						local idx = base + j
						xc._storage[idx] = x._storage[idx] - mean
					end
				end
				task.synchronize()
			end
			local inv: Types.Tensor = Tensor.zeros({ D, 1 })
			local y = Tensor.zeros(x._shape, x._dtype, x._requiresGrad)
			task.desynchronize()
			for i = 1, D do
				local invVal = 1 / math.sqrt(varUsed._storage[i] + (eps :: number))
				inv._storage[i] = invVal
				local base = (i - 1) * B
				local g = gamma._storage[i]
				local b = beta._storage[i]
				for j = 1, B do
					local idx = base + j
					y._storage[idx] = g * xc._storage[idx] * invVal + b
				end
			end
			task.synchronize()
			return Tape.makeNode(y, { x, gamma, beta }, function(gy: Types.Tensor)
				if gamma._requiresGrad then
					task.desynchronize()
					if not gamma._grad then
						gamma._grad = Tensor.zeros({ D, 1 })
					end
					for i = 1, D do
						local acc = 0
						for j = 1, B do
							local idx = (i - 1) * B + j
							acc += gy._storage[idx] * xc._storage[idx] * inv._storage[i]
						end
						gamma._grad._storage[i] += acc
					end
					task.synchronize()
				end
				if beta._requiresGrad then
					task.desynchronize()
					if not beta._grad then
						beta._grad = Tensor.zeros({ D, 1 })
					end
					for i = 1, D do
						local acc = 0
						for j = 1, B do
							local idx = (i - 1) * B + j
							acc += gy._storage[idx]
						end
						beta._grad._storage[i] += acc
					end
					task.synchronize()
				end
				if x._requiresGrad then
					task.desynchronize()
					if not x._grad then
						x._grad = Tensor.zeros(x._shape)
					end
					local xGrad = x._grad :: Types.Tensor
					local sumGy = table.create(D, 0)
					local sumGyXhat = table.create(D, 0)
					for i = 1, D do
						local base = (i - 1) * B
						local invI = inv._storage[i]
						local accGy = 0
						local accGyXhat = 0
						for j = 1, B do
							local idx = base + j
							local gyVal = gy._storage[idx]
							accGy += gyVal
							accGyXhat += gyVal * (xc._storage[idx] * invI)
						end
						sumGy[i] = accGy
						sumGyXhat[i] = accGyXhat
					end
					local M = B
					for i = 1, D do
						local base = (i - 1) * B
						local g = gamma._storage[i]
						local invI = inv._storage[i]
						local sumGyI = sumGy[i]
						local sumGyXhatI = sumGyXhat[i]
						for j = 1, B do
							local idx = base + j
							local gyVal = gy._storage[idx]
							local xhat = xc._storage[idx] * invI
							local term = (M * gyVal - sumGyI - xhat * sumGyXhatI) / M
							xGrad._storage[idx] += g * invI * term
						end
					end
					task.synchronize()
				end
			end)
		end,
		parameters = function(_)
			return { gamma, beta }
		end,
	}
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")
local Util = require("../Util")
local assert = Util.Assert

--@param C number
--@return table
return function(C: number): Types.Module<Types.Tensor, Types.Tensor>
	local gamma = Tensor.fromArray(table.create(C, 1), { C, 1 }, "f64", true)
	local beta = Tensor.zeros({ C, 1 }, "f64", true)
	local self
	self = {
		forward = function(_, x: Types.Tensor): Types.Tensor
			local Cx, H, W, B = x._shape[1], x._shape[2], x._shape[3], x._shape[4]
			assert(Cx == C, "BN2d C mismatch")
			task.desynchronize()
			local M = H * W * B
			local y = Tensor.zeros(
				{ C, H, W, B },
				x._dtype,
				(x._requiresGrad or gamma._requiresGrad or beta._requiresGrad) == true
			)
			local mu = table.create(C, 0)
			local inv = table.create(C, 0)
			for c = 1, C do
				local sum = 0
				for i = 1, H do
					for j = 1, W do
						for b = 1, B do
							local xi = (((c - 1) * H + (i - 1)) * W + (j - 1)) * B + b
							sum += x._storage[xi]
						end
					end
				end
				mu[c] = sum / M
			end
			local eps = 1e-5
			for c = 1, C do
				local v = 0
				for i = 1, H do
					for j = 1, W do
						for b = 1, B do
							local xi = (((c - 1) * H + (i - 1)) * W + (j - 1)) * B + b
							local d = x._storage[xi] - mu[c]
							v += d * d
						end
					end
				end
				inv[c] = 1 / math.sqrt(v / M + eps)
			end
			-- y = gamma * (x - mu) * inv + beta
			for c = 1, C do
				local g = gamma._storage[c]
				local bt = beta._storage[c]
				local muc = mu[c]
				local ic = inv[c]
				for i = 1, H do
					for j = 1, W do
						for b = 1, B do
							local xi = (((c - 1) * H + (i - 1)) * W + (j - 1)) * B + b
							y._storage[xi] = g * (x._storage[xi] - muc) * ic + bt
						end
					end
				end
			end
			task.synchronize()
			return Tape.makeNode(y, { x, gamma, beta }, function(gy: Types.Tensor)
				local M = H * W * B
				-- dBeta, dGamma
				if beta._requiresGrad or gamma._requiresGrad then
					task.desynchronize()
					if beta._requiresGrad then
						beta._grad = beta._grad or Tensor.zeros({ C, 1 }, beta._dtype, false)
						for c = 1, C do
							local acc = 0
							for i = 1, H do
								for j = 1, W do
									for b = 1, B do
										local yi = (((c - 1) * H + (i - 1)) * W + (j - 1)) * B + b
										acc += gy._storage[yi]
									end
								end
							end
							beta._grad._storage[c] += acc
						end
					end
					if gamma._requiresGrad then
						gamma._grad = gamma._grad or Tensor.zeros({ C, 1 }, gamma._dtype, false)
						for c = 1, C do
							local acc = 0
							local muc, ic = mu[c], inv[c]
							for i = 1, H do
								for j = 1, W do
									for b = 1, B do
										local xi = (((c - 1) * H + (i - 1)) * W + (j - 1)) * B + b
										local xhat = (x._storage[xi] - muc) * ic
										local yi = xi
										acc += gy._storage[yi] * xhat
									end
								end
							end
							gamma._grad._storage[c] += acc
						end
					end
					task.synchronize()
				end
				-- dX
				if x._requiresGrad then
					x._grad = x._grad or Tensor.zeros(x._shape, x._dtype, false)
					task.desynchronize()
					for c = 1, C do
						local muc, ic = mu[c], inv[c]
						local g = gamma._storage[c]
						local sumGy, sumGyXhat = 0, 0
						for i = 1, H do
							for j = 1, W do
								for b = 1, B do
									local xi = (((c - 1) * H + (i - 1)) * W + (j - 1)) * B + b
									local yi = xi
									local gyij = gy._storage[yi]
									sumGy += gyij
									sumGyXhat += gyij * ((x._storage[xi] - muc) * ic)
								end
							end
						end
						for i = 1, H do
							for j = 1, W do
								for b = 1, B do
									local xi = (((c - 1) * H + (i - 1)) * W + (j - 1)) * B + b
									local gyij = gy._storage[xi]
									local term = (M * gyij - sumGy - ((x._storage[xi] - muc) * ic) * sumGyXhat) / M
									x._grad._storage[xi] += g * ic * term
								end
							end
						end
					end
					task.synchronize()
				end
			end)
		end,
		parameters = function()
			return { gamma, beta }
		end,
	}
	return self
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local convop = require("../ops/Conv2d")
local Types = require("../Types")

-- {C_in,H,W,B} -> {C_out,H-KH+1,W-KW+1,B}, stride=1, pad=0 (for now), with bias
return function(C_in: number, C_out: number, KH: number, KW: number): Types.Module<Types.Tensor, Types.Tensor>
	local W = Tensor.zeros({ C_out, C_in, KH, KW }, "f64", true)
	local b = Tensor.zeros({ C_out, 1 }, "f64", true)
	local self = {}
	self.parameters = function(_)
		return { W, b }
	end
	self.train = function(_) end
	self.forward = function(_, x: Types.Tensor): Types.Tensor
		local out = convop(x, W)
		-- add bias
		task.desynchronize()
		local F, Ho, Wo, B = out._shape[1], out._shape[2], out._shape[3], out._shape[4]
		for f = 1, F do
			local bf = b._storage[f]
			for i = 1, Ho do
				for j = 1, Wo do
					for bb = 1, B do
						local yi = (((f - 1) * Ho + (i - 1)) * Wo + (j - 1)) * B + bb
						out._storage[yi] += bf
					end
				end
			end
		end
		task.synchronize()
		local old_backprop = out._backprop
		local parents = out._parents and table.clone(out._parents) or {}
		table.insert(parents, b)
		if out._requiresGrad or b._requiresGrad then
			return Tape.makeNode(out, parents, function(gy: Types.Tensor)
				if b._requiresGrad then
					b._grad = b._grad or Tensor.zeros(b._shape, b._dtype, false)
					task.desynchronize()
					for f = 1, F do
						local acc = 0
						for i = 1, Ho do
							for j = 1, Wo do
								for bb = 1, B do
									local yi = (((f - 1) * Ho + (i - 1)) * Wo + (j - 1)) * B + bb
									acc += gy._storage[yi]
								end
							end
						end
						b._grad._storage[f] += acc
					end
					task.synchronize()
				end
				if old_backprop then
					old_backprop(gy)
				end
			end)
		end
		return out
	end
	return self
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local convT = require("../ops/ConvTranspose2d")
local Types = require("../Types")

--{C_in,H,W,B} -> {C_out, H+KH-1, W+KW-1, B}
return function(C_in: number, C_out: number, KH: number, KW: number): Types.Module<Types.Tensor, Types.Tensor>
	local W = Tensor.zeros({ C_in, C_out, KH, KW }, "f64", true)
	local b = Tensor.zeros({ C_out, 1 }, "f64", true)
	local self = {}
	self.parameters = function(_)
		return { W, b }
	end
	self.train = function(_) end
	self.forward = function(_, x: Types.Tensor): Types.Tensor
		local out = convT(x, W)
		task.desynchronize()
		local F, Ho, Wo, B = out._shape[1], out._shape[2], out._shape[3], out._shape[4]
		local os = out._storage
		local bs = b._storage
		local WoB = Wo * B
		for f = 1, F do
			local bf = bs[f]
			local f_offset = (f - 1) * Ho * Wo * B
			for i = 1, Ho do
				local i_offset = (i - 1) * WoB
				for j = 1, Wo do
					local j_offset = (j - 1) * B
					for bb = 1, B do
						local idx = f_offset + i_offset + j_offset + bb
						os[idx] += bf
					end
				end
			end
		end
		task.synchronize()
		local old_backprop = out._backprop
		local parents = out._parents and table.clone(out._parents) or {}
		table.insert(parents, b)
		if out._requiresGrad or b._requiresGrad then
			return Tape.makeNode(out, parents, function(gy: Types.Tensor)
				if b._requiresGrad then
					b._grad = b._grad or Tensor.zeros(b._shape, b._dtype, false)
					task.desynchronize()
					local gbs = b._grad._storage
					local gys = gy._storage
					for f = 1, F do
						local acc = 0
						local f_offset = (f - 1) * Ho * Wo * B
						for i = 1, Ho do
							local i_offset = (i - 1) * WoB
							for j = 1, Wo do
								local j_offset = (j - 1) * B
								for bb = 1, B do
									local idx = f_offset + i_offset + j_offset + bb
									acc += gys[idx]
								end
							end
						end
						gbs[f] += acc
					end
					task.synchronize()
				end
				if old_backprop then
					old_backprop(gy)
				end
			end)
		end
		return out
	end
	return self
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")

return function(p: number?, training: boolean?, rng: Random?): Types.Module<Types.Tensor, Types.Tensor>
	local p = p or 0.2
	local rng = rng or Random.new()
	local keep = 1 - p
	return {
		_train = type(training) == "boolean" and training or true,
		train = function(self, mode)
			self._train = mode == nil and true or mode
		end,
		parameters = function(self)
			return {}
		end,
		forward = function(self, x: Types.Tensor): Types.Tensor
			if not self._train or p <= 0 then
				return x
			end
			local D, B = x._shape[1], x._shape[2]
			local out = Tensor.zeros(x._shape, x._dtype, x._requiresGrad)
			task.desynchronize()
			local smask = table.create(B, 0)
			for b = 1, B do
				smask[b] = (rng:NextNumber(0, 1) < keep) and 1 or 0
			end
			for d = 1, D do
				local base = (d - 1) * B
				for b = 1, B do
					local m = smask[b]
					out._storage[base + b] = x._storage[base + b] * m / keep
				end
			end
			task.synchronize()
			if not out._requiresGrad then
				return out
			end
			return Tape.makeNode(out, { x }, function(gy: Types.Tensor)
				if x._requiresGrad then
					x._grad = x._grad or Tensor.zeros(x._shape, x._dtype, false)
					task.desynchronize()
					for d = 1, D do
						local base = (d - 1) * B
						for b = 1, B do
							x._grad._storage[base + b] += gy._storage[base + b] * smask[b] / keep
						end
					end
					task.synchronize()
				end
			end)
		end,
	}
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")

return function(p: number?, training: boolean?, rng: Random?): Types.Module<Types.Tensor, Types.Tensor>
	local rng = rng or Random.new()
	local drop = p or 0.2
	local keep = 1 - drop
	return {
		_train = type(training) == "boolean" and training or true,
		train = function(self, mode)
			self._train = mode == nil and true or mode
		end,
		parameters = function(self)
			return {}
		end,
		forward = function(self, x: Types.Tensor): Types.Tensor
			if not self._train or drop <= 0 then
				return x
			end
			task.desynchronize()
			local y = Tensor.zeros(x._shape, x._dtype, x._requiresGrad)
			local mask = table.create(#x._storage, 0)
			for i = 1, #x._storage do
				local m = (rng:NextNumber(0, 1) >= drop) and 1 or 0
				mask[i] = m
				y._storage[i] = x._storage[i] * m / keep
			end
			task.synchronize()
			return Tape.makeNode(y, { x }, function(gy: Types.Tensor)
				if x._requiresGrad then
					task.desynchronize()
					if not x._grad then
						x._grad = Tensor.zeros(x._shape)
					end
					for i = 1, #x._grad._storage do
						x._grad._storage[i] += gy._storage[i] * mask[i] / keep
					end
					task.synchronize()
				end
			end)
		end,
	}
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")

return function(numEmbeds: number, dim: number): Types.Module<Types.Tensor, Types.Tensor>
	local W = Tensor.zeros({ dim, numEmbeds }, "f64", true)
	return {
		forward = function(_, idx: { number }): Types.Tensor
			task.desynchronize()
			local B = #idx
			local Y = Tensor.zeros({ dim, B }, "f64", W._requiresGrad)
			for j = 1, B do
				local k = math.clamp(idx[j], 1, numEmbeds)
				for i = 1, dim do
					Y._storage[(i - 1) * B + j] = W._storage[(i - 1) * numEmbeds + k]
				end
			end
			task.synchronize()
			return Tape.makeNode(Y, { W }, function(gY: Types.Tensor)
				if W._requiresGrad then
					task.desynchronize()
					if not W._grad then
						W._grad = Tensor.zeros({ dim, numEmbeds })
					end
					for j = 1, B do
						local k = math.clamp(idx[j], 1, numEmbeds)
						for i = 1, dim do
							W._grad._storage[(i - 1) * numEmbeds + k] += gY._storage[(i - 1) * B + j]
						end
					end
					task.synchronize()
				end
			end)
		end,
		parameters = function(_)
			return { W }
		end,
	}
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")

-- Flattens everything but the last dimension (batch) for NCHW+B -> {C*H*W,B}
return function(): Types.Module<Types.Tensor, Types.Tensor>
	return {
		forward = function(_, x: Types.Tensor): Types.Tensor
			local shape = x._shape
			if #shape == 2 then
				return x
			end
			if #shape == 4 then
				local C, H, W, B = shape[1], shape[2], shape[3], shape[4]
				local out = Tensor.zeros({ C * H * W, B }, x._dtype, x._requiresGrad)
				task.desynchronize()
				for c = 1, C do
					for i = 1, H do
						for j = 1, W do
							for b = 1, B do
								local xi = (((c - 1) * H + (i - 1)) * W + (j - 1)) * B + b
								local yi = ((((c - 1) * H + (i - 1)) * W + (j - 1)) * B) + b
								out._storage[yi] = x._storage[xi]
							end
						end
					end
				end
				task.synchronize()
				if not out._requiresGrad then
					return out
				end
				return Tape.makeNode(out, { x }, function(gy: Types.Tensor)
					if x._requiresGrad then
						x._grad = x._grad or Tensor.zeros(x._shape, x._dtype, false)
						task.desynchronize()
						for c = 1, C do
							for i = 1, H do
								for j = 1, W do
									for b = 1, B do
										local xi = (((c - 1) * H + (i - 1)) * W + (j - 1)) * B + b
										local yi = ((((c - 1) * H + (i - 1)) * W + (j - 1)) * B) + b
										x._grad._storage[xi] += gy._storage[yi]
									end
								end
							end
						end
						task.synchronize()
					end
				end)
			else
				error(`Flatten: unsupported tensor rank {#shape}.`)
			end
		end,
		parameters = function(_)
			return {}
		end,
		train = function(_) end,
	}
end


--!native
--!optimize 2
local Tape = require("../autograd/Tape")
local Types = require("../Types")
local FN = {}

function FN.relu(x: Types.Tensor): Types.Tensor
	return Tape.relu(x)
end

function FN.linear(w: Types.Tensor, b: Types.Tensor, x: Types.Tensor): Types.Tensor
	task.desynchronize()
	local out = Tape.matmul(w, x)
	local B = x._shape[2]
	for i = 1, w._shape[1] do
		local bi = b._storage[i]
		for j = 1, B do
			out._storage[(i - 1) * B + j] += bi
		end
	end
	task.synchronize()
	return out
end

return FN


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Types = require("../Types")
local Linear = require("./Linear")
local Activations = require("./Activations")
local Dropout = require("./Dropout")
local Util = require("../Util")
local assert = Util.Assert

type Tensor = Types.Tensor
type Module = Types.Module

local Fused = {}

--@param x Tensor {inDim, batch}
--@param W Tensor {outDim, inDim}
--@param b Tensor {outDim, 1}
--@returns Tensor {outDim, batch} = ReLU(Wx + b)
function Fused.linearReLU(x: Tensor, W: Tensor, b: Tensor): Tensor
	local inDim = x._shape[1]
	local batch = x._shape[2]
	local outDim = W._shape[1]

	local y = Tensor.zeros({ outDim, batch }, x._dtype, false)
	local yStore = y._storage
	local xStore = x._storage
	local wStore = W._storage
	local bStore = b._storage

	task.desynchronize()
	for i = 1, outDim do
		local bias = bStore[i]
		for j = 1, batch do
			local sum = 0.0
			-- dot(W[i,:], x[:,j])
			for k = 1, inDim do
				local wIdx = (i - 1) * inDim + k -- W row-major
				local xIdx = (k - 1) * batch + j -- x {rows=inDim, batch}
				sum += wStore[wIdx] * xStore[xIdx]
			end

			sum += bias
			if sum < 0 then
				sum = 0
			end

			local yIdx = (i - 1) * batch + j
			yStore[yIdx] = sum
		end
	end
	task.synchronize()

	return y
end

local function gelu_scalar(x: number): number
	-- 0.5 * x * (1 + tanh(√(2/π)*(x + 0.044715 x^3)))
	local c = math.sqrt(2 / math.pi)
	local x3 = x * x * x
	local inner = c * (x + 0.044715 * x3)
	local t = math.tanh(inner)
	return 0.5 * x * (1 + t)
end

--@param x Tensor {inDim, batch}
--@param W Tensor {outDim, inDim}
--@param b Tensor {outDim, 1}
--@returns Tensor {outDim, batch} = GELU(Wx + b)
function Fused.linearGELU(x: Tensor, W: Tensor, b: Tensor): Tensor
	local inDim = x._shape[1]
	local batch = x._shape[2]
	local outDim = W._shape[1]

	local y = Tensor.zeros({ outDim, batch }, x._dtype, false)
	local yStore = y._storage
	local xStore = x._storage
	local wStore = W._storage
	local bStore = b._storage

	task.desynchronize()
	for i = 1, outDim do
		local bias = bStore[i]
		for j = 1, batch do
			local sum = 0.0
			for k = 1, inDim do
				local wIdx = (i - 1) * inDim + k
				local xIdx = (k - 1) * batch + j
				sum += wStore[wIdx] * xStore[xIdx]
			end

			sum += bias
			sum = gelu_scalar(sum)

			local yIdx = (i - 1) * batch + j
			yStore[yIdx] = sum
		end
	end
	task.synchronize()

	return y
end

--@param x Tensor {inDim, batch}
--@param W Tensor {outDim, inDim}
--@param b Tensor {outDim, 1}
--@param p number dropout probability in [0,1)
--@param rng (() -> number)? (optional) → number in [0,1), default math.random
--@returns { Tensor, Tensor }
function Fused.linearDropoutReLU(x: Tensor, W: Tensor, b: Tensor, p: number, rng: (() -> number)?): (Tensor, Tensor)
	local rng = rng or math.random

	local keep = 1.0 - p
	assert(keep > 0, "Fused.linearDropoutReLU: keep prob must be > 0")

	local inDim = x._shape[1]
	local batch = x._shape[2]
	local outDim = W._shape[1]

	local y = Tensor.zeros({ outDim, batch }, x._dtype, false)
	local mask = Tensor.zeros({ outDim, batch }, x._dtype, false)

	local yStore = y._storage
	local mStore = mask._storage
	local xStore = x._storage
	local wStore = W._storage
	local bStore = b._storage

	local scale = 1.0 / keep

	task.desynchronize()
	for i = 1, outDim do
		local bias = bStore[i]
		for j = 1, batch do
			local sum = 0.0
			for k = 1, inDim do
				local wIdx = (i - 1) * inDim + k
				local xIdx = (k - 1) * batch + j
				sum += wStore[wIdx] * xStore[xIdx]
			end

			sum += bias

			local idx = (i - 1) * batch + j
			local r = rng()
			if r < keep and sum > 0 then
				mStore[idx] = 1.0
				yStore[idx] = sum * scale
			else
				mStore[idx] = 0.0
				yStore[idx] = 0.0
			end
		end
	end
	task.synchronize()
	return y, mask
end

export type LinearReLU = {
	_linear: Module,
	parameters: (self: LinearReLU) -> { Types.Tensor },
	forward: (self: LinearReLU, x: Tensor) -> Tensor,
}

export type LinearGELU = {
	_linear: Module,
	parameters: (self: LinearGELU) -> { Types.Tensor },
	forward: (self: LinearGELU, x: Tensor) -> Tensor,
}

export type LinearDropoutReLU = {
	_linear: Module,
	_drop: Module,
	parameters: (self: LinearDropoutReLU) -> { Types.Tensor },
	forward: (self: LinearDropoutReLU, x: Tensor) -> Tensor,
}

local Modules = {}

local function linearReLUParameters(self: LinearReLU)
	return self._linear:parameters()
end

local function linearReLUForward(self: LinearReLU, x: Tensor): Tensor
	local z = self._linear:forward(x)
	return Activations.ReLU(z)
end

function Modules.LinearReLU(inDim: number, outDim: number, init: ((number, number) -> number)?): Module
	local linear = Linear(inDim, outDim, init)
	local self: LinearReLU = {
		_linear = linear,
		parameters = linearReLUParameters,
		forward = linearReLUForward,
	}
	return (self :: any) :: Module
end

local function linearGELUParameters(self: LinearGELU)
	return self._linear:parameters()
end

local function linearGELUForward(self: LinearGELU, x: Tensor): Tensor
	local z = self._linear:forward(x)
	return Activations.GELU(z)
end

function Modules.LinearGELU(inDim: number, outDim: number, init: ((number, number) -> number)?): Module
	local linear = Linear(inDim, outDim, init)
	local self: LinearGELU = {
		_linear = linear,
		parameters = linearGELUParameters,
		forward = linearGELUForward,
	}
	return (self :: any) :: Module
end

local function linearDropoutReLUParameters(self: LinearDropoutReLU)
	return self._linear:parameters()
end

local function linearDropoutReLUForward(self: LinearDropoutReLU, x: Tensor): Tensor
	local z = self._linear:forward(x)
	local d = self._drop:forward(z)
	return Activations.ReLU(d)
end

function Modules.LinearDropoutReLU(
	inDim: number,
	outDim: number,
	p: number,
	init: ((number, number) -> number)?
): Module
	local linear = Linear(inDim, outDim, init)
	local drop = Dropout(p)
	local self: LinearDropoutReLU = {
		_linear = linear,
		_drop = drop,
		parameters = linearDropoutReLUParameters,
		forward = linearDropoutReLUForward,
	}
	return (self :: any) :: Module
end

Fused.Modules = Modules

return Fused


--!native
--!optimize 2
local Linear = require("./Linear")
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")

return function(inF: number, hidF: number): Types.Module<Types.Tensor, Types.Tensor>
	local Wx = Linear(inF, 3 * hidF)
	local Wh = Linear(hidF, 3 * hidF)
	local self
	self = {
		forward = function(_, x: Types.Tensor, h: Types.Tensor): Types.Tensor
			task.desynchronize()
			local gx = Wx:forward(x) -- {3H,B}
			local gh = Wh:forward(h) -- {3H,B}
			local B = x._shape[2]
			local z = Tensor.zeros({ hidF, B }, gx._dtype, gx._requiresGrad or gh._requiresGrad)
			local r = Tensor.zeros({ hidF, B }, gx._dtype, gx._requiresGrad or gh._requiresGrad)
			local n = Tensor.zeros({ hidF, B }, gx._dtype, gx._requiresGrad or gh._requiresGrad)
			local y = Tensor.zeros({ hidF, B }, gx._dtype, gx._requiresGrad or gh._requiresGrad)
			for i = 1, hidF do
				local zi = i -- z rows
				local ri = hidF + i -- r rows
				local ni = 2 * hidF + i -- n rows
				for b = 1, B do
					local zb = (zi - 1) * B + b
					local rb = (ri - 1) * B + b
					local nb = (ni - 1) * B + b
					local zpre = gx._storage[zb] + gh._storage[zb]
					local rpre = gx._storage[rb] + gh._storage[rb]
					local zv = 1 / (1 + math.exp(-zpre))
					local rv = 1 / (1 + math.exp(-rpre))
					z._storage[zb] = zv
					r._storage[rb] = rv
					local npre = gx._storage[nb] + rv * gh._storage[nb]
					local nv = math.tanh(npre)
					n._storage[nb] = nv
					y._storage[(i - 1) * B + b] = (1 - zv) * nv + zv * h._storage[(i - 1) * B + b]
				end
			end
			task.synchronize()
			return Tape.makeNode(y, { gx, gh, h }, function(gy: Types.Tensor)
				local B = y._shape[2]
				if gx._requiresGrad then
					gx._grad = gx._grad or Tensor.zeros(gx._shape, gx._dtype, false)
				end
				if gh._requiresGrad then
					gh._grad = gh._grad or Tensor.zeros(gh._shape, gh._dtype, false)
				end
				if h._requiresGrad then
					h._grad = h._grad or Tensor.zeros(h._shape, h._dtype, false)
				end
				task.desynchronize()
				for i = 1, hidF do
					local zi = i
					local ri = hidF + i
					local ni = 2 * hidF + i
					for b = 1, B do
						local hb = (i - 1) * B + b
						local zb = (zi - 1) * B + b
						local rb = (ri - 1) * B + b
						local nb = (ni - 1) * B + b
						local zv = z._storage[zb]
						local rv = r._storage[rb]
						local nv = n._storage[nb]
						local gyv = gy._storage[hb]
						-- y = (1 - z) * n + z * h
						local d_n = gyv * (1 - zv)
						local d_z = gyv * (h._storage[hb] - nv)
						if h._requiresGrad then
							h._grad._storage[hb] += gyv * zv
						end
						-- n = tanh(gx_n + r * gh_n)
						local dnpre = d_n * (1 - nv * nv)
						local d_gx_n = dnpre
						local d_gh_n = dnpre * rv
						local d_r = dnpre * gh._storage[nb]
						-- z,r sigmoids
						local dpre_z = d_z * zv * (1 - zv)
						local dpre_r = d_r * rv * (1 - rv)
						-- accumulate to gx/gh slices
						if gx._requiresGrad then
							gx._grad._storage[zb] += dpre_z
							gx._grad._storage[rb] += dpre_r
							gx._grad._storage[nb] += d_gx_n
						end
						if gh._requiresGrad then
							gh._grad._storage[zb] += dpre_z
							gh._grad._storage[rb] += dpre_r
							gh._grad._storage[nb] += d_gh_n
						end
					end
				end
				task.synchronize()
			end)
		end,
		parameters = function(_): { Types.Tensor }
			local p = Wx:parameters()
			for _, t in Wh:parameters() do
				table.insert(p, t)
			end
			return p
		end,
	}
	return self
end


--!native
--!optimize 2
local Linear = require("./Linear")
local A = require("./Activations")
local Types = require("../Types")

--@class GatedMLP
--@field fc_in any -- Linear(in, 2H)
--@field fc_out any -- Linear(H, out)

local function forward(self, x: Types.Tensor): Types.Tensor
	local twoH = self.fc_in:forward(x) -- {2H, B}
	local h = A.SwiGLUSplit(twoH) -- {H, B}
	return self.fc_out:forward(h) -- {O, B}
end

local function parameters(self): { Types.Tensor }
	local p, a, b = {}, self.fc_in:parameters(), self.fc_out:parameters()
	for i = 1, #a do
		p[#p + 1] = a[i]
	end
	for i = 1, #b do
		p[#p + 1] = b[i]
	end
	return p
end

local function new(inDim: number, hidden: number, outDim: number)
	return {
		fc_in = Linear(inDim, 2 * hidden),
		fc_out = Linear(hidden, outDim),
		forward = forward,
		parameters = parameters,
	}
end

return new


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")
local Util = require("../Util")
local assert = Util.Assert

return function(num_groups: number, num_channels: number, eps: number?): Types.Module<Types.Tensor, Types.Tensor>
	eps = eps or 1e-5
	assert(num_channels % num_groups == 0, "GroupNorm: num_channels must be divisible by num_groups")

	local gamma = Tensor.ones({ num_channels, 1 }, "f64", true)
	local beta = Tensor.zeros({ num_channels, 1 }, "f64", true)

	local self
	self = {
		forward = function(_, x: Types.Tensor): Types.Tensor
			local C = x._shape[1]
			assert(C == num_channels, "GroupNorm: channel mismatch")
			local shape = x._shape
			local rank = #shape
			local B = shape[rank]
			local spatial_size = 1
			for i = 2, rank - 1 do
				spatial_size *= shape[i]
			end
			local G = num_groups
			local C_per_G = C / G
			local group_size = C_per_G * spatial_size
			local xs = x._storage
			task.desynchronize()
			local stats = table.create(G * B * 2, 0)
			local stride_C = Util.sizeFromShape(shape) / C
			for c = 1, C do
				local g = math.ceil(c / C_per_G)
				local g_idx = (g - 1) * B
				local c_base = (c - 1) * stride_C
				for i = 0, stride_C - 1 do
					local b = (i % B) + 1
					local val = xs[c_base + i + 1]
					stats[g_idx + b] += val
				end
			end
			local inv_size = 1 / group_size
			for k = 1, G * B do
				stats[k] *= inv_size
			end
			for c = 1, C do
				local g = math.ceil(c / C_per_G)
				local g_idx = (g - 1) * B
				local c_base = (c - 1) * stride_C
				for i = 0, stride_C - 1 do
					local b = (i % B) + 1
					local val = xs[c_base + i + 1]
					local m = stats[g_idx + b]
					stats[G * B + g_idx + b] += (val - m) ^ 2
				end
			end
			for k = G * B + 1, G * B * 2 do
				stats[k] = math.sqrt(stats[k] * inv_size + eps)
			end
			local y = Tensor.zeros(shape, x._dtype, x._requiresGrad)
			local ys = y._storage
			local gs = gamma._storage
			local bs = beta._storage
			for c = 1, C do
				local g = math.ceil(c / C_per_G)
				local g_idx = (g - 1) * B
				local c_base = (c - 1) * stride_C
				local gam = gs[c]
				local bet = bs[c]
				for i = 0, stride_C - 1 do
					local b = (i % B) + 1
					local val = xs[c_base + i + 1]
					local m = stats[g_idx + b]
					local s = stats[G * B + g_idx + b]
					-- y = gamma * (x - mean) / std + beta
					ys[c_base + i + 1] = gam * (val - m) / s + bet
				end
			end
			task.synchronize()
			return Tape.makeNode(y, { x, gamma, beta }, function(gy: Types.Tensor)
				if x._requiresGrad then
					if not x._grad then
						x._grad = Tensor.zeros(shape, x._dtype, false)
					end
				end
				if gamma._requiresGrad and not gamma._grad then
					gamma._grad = Tensor.zeros({ C, 1 }, "f64", false)
				end
				if beta._requiresGrad and not beta._grad then
					beta._grad = Tensor.zeros({ C, 1 }, "f64", false)
				end
				task.desynchronize()
				local gys = gy._storage
				local dxs = x._requiresGrad and (x._grad :: Types.Tensor)._storage
				local dgs = gamma._requiresGrad and gamma._grad._storage
				local dbs = beta._requiresGrad and beta._grad._storage
				local back_stats = table.create(G * B * 2, 0)
				for c = 1, C do
					local g = math.ceil(c / C_per_G)
					local g_idx = (g - 1) * B
					local c_base = (c - 1) * stride_C
					local gam = gs[c]
					for i = 0, stride_C - 1 do
						local b = (i % B) + 1
						local idx = c_base + i + 1
						local gy_val = gys[idx]
						local val = xs[idx]
						local m = stats[g_idx + b]
						local s = stats[G * B + g_idx + b]
						local x_hat = (val - m) / s
						if dgs then
							dgs[c] += gy_val * x_hat
						end
						if dbs then
							dbs[c] += gy_val
						end
						local dy_gamma = gy_val * gam
						back_stats[g_idx + b] += dy_gamma
						back_stats[G * B + g_idx + b] += dy_gamma * x_hat
					end
				end
				if dxs then
					local N = group_size
					local dxs_nonil = dxs :: { number }
					for c = 1, C do
						local g = math.ceil(c / C_per_G)
						local g_idx = (g - 1) * B
						local c_base = (c - 1) * stride_C
						local gam = gs[c]
						for i = 0, stride_C - 1 do
							local b = (i % B) + 1
							local idx = c_base + i + 1
							local m = stats[g_idx + b]
							local s = stats[G * B + g_idx + b]
							local val = xs[idx]
							local x_hat = (val - m) / s
							local sum_dy_gamma = back_stats[g_idx + b]
							local sum_dy_gamma_xhat = back_stats[G * B + g_idx + b]
							local dy = gys[idx]
							-- dL/dx = (1/s) * ( dy*gamma - sum_dy_gamma/N - x_hat * sum_dy_gamma_xhat/N )
							local term = (dy * gam) - (sum_dy_gamma / N) - (x_hat * sum_dy_gamma_xhat / N)
							dxs_nonil[idx] += term / s
						end
					end
				end
				task.synchronize()
			end)
		end,
		parameters = function(_)
			return { gamma, beta }
		end,
	}
	return self
end


--!native
--!optimize 2
local Linear = require("./Linear")
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")

return function(inF: number, hidF: number): Types.Module<Types.Tensor, Types.Tensor>
	local Wx = Linear(inF, 4 * hidF)
	local Wh = Linear(hidF, 4 * hidF)
	local self
	self = {
		forward = function(_, x: Types.Tensor, h: Types.Tensor, c: Types.Tensor): Types.Tensor
			task.desynchronize()
			local B = x._shape[2]
			local gx = Wx:forward(x) -- {4H,B}
			local gh = Wh:forward(h) -- {4H,B}
			local i = Tensor.zeros({ hidF, B }, gx._dtype, gx._requiresGrad or gh._requiresGrad)
			local f = Tensor.zeros({ hidF, B }, gx._dtype, gx._requiresGrad or gh._requiresGrad)
			local g = Tensor.zeros({ hidF, B }, gx._dtype, gx._requiresGrad or gh._requiresGrad)
			local o = Tensor.zeros({ hidF, B }, gx._dtype, gx._requiresGrad or gh._requiresGrad)
			local cNew = Tensor.zeros(
				{ hidF, B },
				gx._dtype,
				(gx._requiresGrad or gh._requiresGrad or (c and c._requiresGrad)) == true
			)
			local hNew = Tensor.zeros(
				{ hidF, B },
				gx._dtype,
				(gx._requiresGrad or gh._requiresGrad or (c and c._requiresGrad)) == true
			)
			for j = 1, hidF do
				local ii = j
				local fi = hidF + j
				local gi = 2 * hidF + j
				local oi = 3 * hidF + j
				for b = 1, B do
					local ib = (ii - 1) * B + b
					local fb = (fi - 1) * B + b
					local gb = (gi - 1) * B + b
					local ob = (oi - 1) * B + b
					local pre_i = gx._storage[ib] + gh._storage[ib]
					local pre_f = gx._storage[fb] + gh._storage[fb]
					local pre_g = gx._storage[gb] + gh._storage[gb]
					local pre_o = gx._storage[ob] + gh._storage[ob]
					local iv = 1 / (1 + math.exp(-pre_i))
					local fv = 1 / (1 + math.exp(-pre_f))
					local gv = math.tanh(pre_g)
					local ov = 1 / (1 + math.exp(-pre_o))
					i._storage[ib] = iv
					f._storage[fb] = fv
					g._storage[gb] = gv
					o._storage[ob] = ov
					local cb = (j - 1) * B + b
					local cprev = c._storage[cb]
					local cnow = fv * cprev + iv * gv
					cNew._storage[cb] = cnow
					hNew._storage[cb] = ov * math.tanh(cnow)
				end
			end
			task.synchronize()
			local hNode = Tape.makeNode(hNew, { gx, gh, c }, function(gyh)
				local B = hNew._shape[2]
				if gx._requiresGrad then
					gx._grad = gx._grad or Tensor.zeros(gx._shape, gx._dtype, false)
				end
				if gh._requiresGrad then
					gh._grad = gh._grad or Tensor.zeros(gh._shape, gh._dtype, false)
				end
				if c._requiresGrad then
					c._grad = c._grad or Tensor.zeros(c._shape, c._dtype, false)
				end
				task.desynchronize()
				for j = 1, hidF do
					local ii = j
					local fi = hidF + j
					local gi = 2 * hidF + j
					local oi = 3 * hidF + j
					for b = 1, B do
						local ib = (ii - 1) * B + b
						local fb = (fi - 1) * B + b
						local gb = (gi - 1) * B + b
						local ob = (oi - 1) * B + b
						local cb = (j - 1) * B + b
						local cnow = cNew._storage[cb]
						local tanh_c = math.tanh(cnow)
						local gyhv = gyh._storage[cb]
						-- h' = o * tanh(c')
						local d_o = gyhv * tanh_c
						local d_c_from_h = gyhv * o._storage[ob] * (1 - tanh_c * tanh_c)
						local dpre_o = d_o * o._storage[ob] * (1 - o._storage[ob])
						if gx._requiresGrad then
							gx._grad._storage[ob] += dpre_o
						end
						if gh._requiresGrad then
							gh._grad._storage[ob] += dpre_o
						end
						if c._requiresGrad then
							local dct = d_c_from_h
							-- c' = f*c + i*g
							local d_i = dct * g._storage[gb]
							local d_f = dct * c._storage[cb]
							local d_g = dct * i._storage[ib]

							local dpre_i = d_i * i._storage[ib] * (1 - i._storage[ib])
							local dpre_f = d_f * f._storage[fb] * (1 - f._storage[fb])
							local dpre_g = d_g * (1 - g._storage[gb] * g._storage[gb])
							if gx._requiresGrad then
								gx._grad._storage[ib] += dpre_i
								gx._grad._storage[fb] += dpre_f
								gx._grad._storage[gb] += dpre_g
							end
							if gh._requiresGrad then
								gh._grad._storage[ib] += dpre_i
								gh._grad._storage[fb] += dpre_f
								gh._grad._storage[gb] += dpre_g
							end
							c._grad._storage[cb] += dct * f._storage[fb]
						end
					end
				end
				task.synchronize()
			end)

			local cNode = Tape.makeNode(cNew, { gx, gh, c }, function(gyc)
				local B = cNew._shape[2]
				if gx._requiresGrad then
					gx._grad = gx._grad or Tensor.zeros(gx._shape, gx._dtype, false)
				end
				if gh._requiresGrad then
					gh._grad = gh._grad or Tensor.zeros(gh._shape, gh._dtype, false)
				end
				if c._requiresGrad then
					c._grad = c._grad or Tensor.zeros(c._shape, c._dtype, false)
				end
				task.desynchronize()
				for j = 1, hidF do
					local ii = j
					local fi = hidF + j
					local gi = 2 * hidF + j
					for b = 1, B do
						local ib = (ii - 1) * B + b
						local fb = (fi - 1) * B + b
						local gb = (gi - 1) * B + b
						local cb = (j - 1) * B + b
						local dct = gyc._storage[cb]
						-- c' = f*c + i*g
						local d_i = dct * g._storage[gb]
						local d_f = dct * c._storage[cb]
						local d_g = dct * i._storage[ib]

						local dpre_i = d_i * i._storage[ib] * (1 - i._storage[ib])
						local dpre_f = d_f * f._storage[fb] * (1 - f._storage[fb])
						local dpre_g = d_g * (1 - g._storage[gb] * g._storage[gb])
						if gx._requiresGrad then
							gx._grad._storage[ib] += dpre_i
							gx._grad._storage[fb] += dpre_f
							gx._grad._storage[gb] += dpre_g
						end
						if gh._requiresGrad then
							gh._grad._storage[ib] += dpre_i
							gh._grad._storage[fb] += dpre_f
							gh._grad._storage[gb] += dpre_g
						end
						if c._requiresGrad then
							c._grad._storage[cb] += dct * f._storage[fb]
						end
					end
				end
				task.synchronize()
			end)
			return hNode, cNode
		end,
		parameters = function(_)
			local p = Wx:parameters()
			for _, t in Wh:parameters() do
				table.insert(p, t)
			end
			return p
		end,
	}
	return self
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Reduce = require("../ops/Reduce")
local Types = require("../Types")
local Util = require("../Util")
local assert = Util.Assert

-- LayerNorm across features for each column (shape {D,B}).
return function(D: number, eps: number?): Types.Module<Types.Tensor, Types.Tensor>
	eps = eps or 1e-5
	local gamma = Tensor.fromArray(table.create(D, 1), { D, 1 }, "f64", true)
	local beta = Tensor.zeros({ D, 1 }, "f64", true)
	local self
	self = {
		forward = function(_, x: Types.Tensor): Types.Tensor
			local D_, B = x._shape[1], x._shape[2]
			assert(D_ == D, "LayerNorm D mismatch")
			task.desynchronize()
			local mu = Reduce.meanCols(x)
			local xc = Tensor.zeros(x._shape, x._dtype, x._requiresGrad)
			for j = 1, B do
				local m = mu._storage[j]
				for i = 1, D do
					xc._storage[(i - 1) * B + j] = x._storage[(i - 1) * B + j] - m
				end
			end
			task.synchronize()
			local var = Reduce.meanCols((function()
				task.desynchronize()
				local sq = Tensor.zeros(x._shape)
				for k = 1, #xc._storage do
					sq._storage[k] = xc._storage[k] ^ 2
				end
				task.synchronize()
				return sq
			end)())
			task.desynchronize()
			local norm = Tensor.zeros(x._shape, x._dtype, x._requiresGrad)
			for j = 1, B do
				local inv = 1 / math.sqrt(var._storage[j] + (eps :: number))
				for i = 1, D do
					norm._storage[(i - 1) * B + j] = xc._storage[(i - 1) * B + j] * inv
				end
			end
			local y = Tensor.zeros(x._shape, x._dtype, x._requiresGrad)
			for j = 1, B do
				for i = 1, D do
					local idx = (i - 1) * B + j
					y._storage[idx] = gamma._storage[i] * norm._storage[idx] + beta._storage[i]
				end
			end
			task.synchronize()
			return Tape.makeNode(y, { x, gamma, beta }, function(gy: Types.Tensor)
				if gamma._requiresGrad then
					task.desynchronize()
					if not gamma._grad then
						gamma._grad = Tensor.zeros({ D, 1 })
					end
					for i = 1, D do
						local acc = 0
						for j = 1, B do
							local idx = (i - 1) * B + j
							acc += gy._storage[idx] * norm._storage[idx]
						end
						gamma._grad._storage[i] += acc
					end
					task.synchronize()
				end
				if beta._requiresGrad then
					task.desynchronize()
					if not beta._grad then
						beta._grad = Tensor.zeros({ D, 1 })
					end
					for i = 1, D do
						local acc = 0
						for j = 1, B do
							local idx = (i - 1) * B + j
							acc += gy._storage[idx]
						end
						beta._grad._storage[i] += acc
					end
					task.synchronize()
				end
				if x._requiresGrad then
					task.desynchronize()
					if not x._grad then
						x._grad = Tensor.zeros(x._shape)
					end
					local xGrad = x._grad :: Types.Tensor
					local inv = table.create(B, 0)
					for j = 1, B do
						inv[j] = 1 / math.sqrt(var._storage[j] + (eps :: number))
					end
					for j = 1, B do
						local invJ = inv[j]
						local sumDxhat = 0
						local sumDxhatXhat = 0
						for i = 1, D do
							local idx = (i - 1) * B + j
							local dxhat = gy._storage[idx] * gamma._storage[i]
							sumDxhat += dxhat
							sumDxhatXhat += dxhat * norm._storage[idx]
						end
						for i = 1, D do
							local idx = (i - 1) * B + j
							local dxhat = gy._storage[idx] * gamma._storage[i]
							local term = dxhat - sumDxhat / D - norm._storage[idx] * sumDxhatXhat / D
							xGrad._storage[idx] += term * invJ
						end
					end
					task.synchronize()
				end
			end)
		end,
		parameters = function(_)
			return { gamma, beta }
		end,
	}
	return self
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")

local function ensureGrad(t: Types.Tensor): Types.Tensor?
	if not t._grad then
		t._grad = Tensor.zeros(t._shape, t._dtype, false)
	end
	return t._grad
end

return function(
	inF: number,
	outF: number,
	init: ((number, number) -> number)?
): Types.Module<Types.Tensor, Types.Tensor>
	local W = Tensor.zeros({ outF, inF }, "f64", true)
	local b = Tensor.zeros({ outF, 1 }, "f64", true)
	if type(init) == "function" then
		task.desynchronize()
		local store = W._storage
		local idx = 1
		for i = 1, outF do
			for _ = 1, inF do
				store[idx] = init(inF, outF)
				idx += 1
			end
		end
		task.synchronize()
	end
	return {
		W = W,
		b = b,
		forward = function(_, x: Types.Tensor): Types.Tensor
			-- A(W): {outF, inF}, B(x): {inF, Batch} -> C: {outF, Batch}
			local M, K = outF, inF
			local K2, N = x._shape[1], x._shape[2]
			if K ~= K2 then
				error("Linear: input shape mismatch")
			end
			local out = Tensor.zeros({ M, N }, W._dtype, W._requiresGrad or b._requiresGrad or x._requiresGrad)
			task.desynchronize()
			local W_store = W._storage
			local x_store = x._storage
			local b_store = b._storage
			local out_store = out._storage
			for i = 1, M do
				local bias_val = b_store[i] or 0
				local w_offset = (i - 1) * K
				local out_offset = (i - 1) * N
				for j = 1, N do
					local acc = 0
					for k = 1, K do
						-- W index: w_offset + k
						-- x index: (k-1)*N + j
						local w_val = W_store[w_offset + k] or 0
						local x_val = x_store[(k - 1) * N + j] or 0
						acc += w_val * x_val
					end
					out_store[out_offset + j] = acc + bias_val -- Fuse bias add
				end
			end
			task.synchronize()
			return Tape.makeNode(out, { W, b, x }, function(gy: Types.Tensor)
				-- dW = gy @ x^T
				-- db = sum(gy, dim=1)
				-- dx = W^T @ gy
				local W_grad, b_grad, x_grad
				if W._requiresGrad then
					W_grad = ensureGrad(W)
				end
				if b._requiresGrad then
					b_grad = ensureGrad(b)
				end
				if x._requiresGrad then
					x_grad = ensureGrad(x)
				end
				task.desynchronize()
				local gy_store = gy._storage
				local B = gy._shape[2] -- Batch size (N)
				if b_grad then
					local bg_store = b_grad._storage
					for i = 1, outF do
						local acc = 0
						local offset = (i - 1) * B
						for j = 1, B do
							acc += gy_store[offset + j] or 0
						end
						bg_store[i] += acc
					end
				end
				-- dW = gy @ x.T
				-- gy: {outF, B}, x: {inF, B} -> dW: {outF, inF}
				-- dW[i, k] = sum_j ( gy[i, j] * x[k, j] )  <-- implicit transpose of x
				if W_grad then
					local wg_store = W_grad._storage
					local x_store = x._storage
					for i = 1, outF do
						local gy_offset = (i - 1) * B
						local w_offset = (i - 1) * inF
						for k = 1, inF do
							local acc = 0
							local x_offset = (k - 1) * B
							for j = 1, B do
								-- gy[i, j] * x[k, j]
								local gy_val = gy_store[gy_offset + j] or 0
								local x_val = x_store[x_offset + j] or 0
								acc += gy_val * x_val
							end
							wg_store[w_offset + k] += acc
						end
					end
				end
				-- dx = W.T @ gy
				-- W: {outF, inF}, gy: {outF, B} -> dx: {inF, B}
				-- dx[k, j] = sum_i ( W[i, k] * gy[i, j] )
				if x_grad then
					local xg_store = x_grad._storage
					local W_store = W._storage
					for k = 1, inF do
						local x_offset = (k - 1) * B
						for j = 1, B do
							local acc = 0
							for i = 1, outF do
								-- W[i, k] => (i-1)*inF + k
								-- gy[i, j] => (i-1)*B + j
								local w_val = W_store[(i - 1) * inF + k] or 0
								local gy_val = gy_store[(i - 1) * B + j] or 0
								acc += w_val * gy_val
							end
							xg_store[x_offset + j] += acc
						end
					end
				end
				task.synchronize()
			end)
		end,
		parameters = function(_)
			return { W, b }
		end,
	}
end


--!native
--!optimize 2
local MP = require("../ops/Pool")
local Types = require("../Types")
return function(KH: number, KW: number, stride: number)
	return {
		forward = function(_, x: Types.Tensor): Types.Tensor
			return MP(x, KH, KW, stride)
		end,
		parameters = function(_)
			return {}
		end,
		train = function(_) end,
	}
end


local Tensor = require("../Tensor")
local Base = require("./Linear")
local Tape = require("../autograd/Tape")
local Types = require("../Types")

return function(inF: number, outF: number, sigma: number?, rng: Random?): Types.Module<Types.Tensor, Types.Tensor>
	sigma = sigma or 0.5
	rng = rng or Random.new()
	local base = Base(inF, outF)
	local bsigma = Tensor.zeros({ outF, 1 }, "f64", true)
	local function resample()
		task.desynchronize()
		for i = 1, #bsigma._storage do
			bsigma._storage[i] = rng:NextNumber(-sigma, sigma)
		end
		task.synchronize()
	end
	return {
		forward = function(_, x: Types.Tensor): Types.Tensor
			local y = base:forward(x)
			task.desynchronize()
			local B = x._shape[2]
			for i = 1, outF do
				local bi = bsigma._storage[i]
				for j = 1, B do
					y._storage[(i - 1) * B + j] += bi
				end
			end
			task.synchronize()
			return Tape.makeNode(y, { base.b, bsigma }, function(gy: Types.Tensor)
				local Bsz = y._shape[2]
				if base.b._requiresGrad then
					base.b._grad = base.b._grad or Tensor.zeros(base.b._shape)
					task.desynchronize()
					for i = 1, outF do
						local acc = 0
						for j = 1, Bsz do
							acc += gy._storage[(i - 1) * Bsz + j]
						end
						base.b._grad._storage[i] += acc
					end
					task.synchronize()
				end
				if bsigma._requiresGrad then
					bsigma._grad = bsigma._grad or Tensor.zeros(bsigma._shape)
					task.desynchronize()
					for i = 1, outF do
						local acc = 0
						for j = 1, Bsz do
							acc += gy._storage[(i - 1) * Bsz + j]
						end
						bsigma._grad._storage[i] += acc
					end
					task.synchronize()
				end
			end)
		end,
		parameters = function(_)
			local p = base:parameters()
			table.insert(p, bsigma)
			return p
		end,
		resample = resample,
	}
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")

local function forward(self: { weight: Types.Tensor, eps: number }, x: Types.Tensor): Types.Tensor
	task.desynchronize()
	local D, B = x._shape[1], x._shape[2]
	local y = Tensor.zeros({ D, B }, x._dtype, x._requiresGrad)
	local w = self.weight

	for j = 1, B do
		local sqsum = 0.0
		for i = 1, D do
			local idx = (i - 1) * B + j
			local v = x._storage[idx]
			sqsum += v * v
		end
		local r = math.sqrt(sqsum / D + self.eps) -- denom
		for i = 1, D do
			local idx = (i - 1) * B + j
			y._storage[idx] = x._storage[idx] * w._storage[i] / r
		end
	end
	task.synchronize()

	return Tape.makeNode(y, { x, w }, function(gy: Types.Tensor)
		if x._requiresGrad or w._requiresGrad then
			task.desynchronize()
			if w._requiresGrad and not w._grad then
				w._grad = Tensor.zeros(w._shape)
			end
			if x._requiresGrad and not x._grad then
				x._grad = Tensor.zeros(x._shape)
			end
			for j = 1, B do
				local sqsum = 0.0
				for i = 1, D do
					local idx = (i - 1) * B + j
					local v = x._storage[idx]
					sqsum += v * v
				end
				local r = math.sqrt(sqsum / D + self.eps)
				local r3 = r * r * r

				local inner = 0.0
				for i = 1, D do
					local idx = (i - 1) * B + j
					inner += w._storage[i] * x._storage[idx] * gy._storage[idx]
				end
				for i = 1, D do
					local idx = (i - 1) * B + j
					local xij = x._storage[idx]
					local gij = gy._storage[idx]
					if w._requiresGrad then
						w._grad._storage[i] += (xij / r) * gij
					end
					if x._requiresGrad and x._grad then
						x._grad._storage[idx] += (w._storage[i] / r) * gij - (xij / (D * r3)) * inner
					end
				end
			end
			task.synchronize()
		end
	end)
end

local function parameters(self: { weight: Types.Tensor, eps: number }): { Types.Tensor }
	return { self.weight }
end

local function new(D: number, eps: number?)
	local weight = Tensor.zeros({ D, 1 }, "f32", true)
	task.desynchronize()
	for i = 1, #weight._storage do
		weight._storage[i] = 1
	end
	task.synchronize()
	return {
		weight = weight,
		eps = eps or 1e-5,
		forward = forward,
		parameters = parameters,
	}
end

return new


--!native
--!optimize 2
local Linear = require("./Linear")
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")

return function(inF: number, hidF: number): Types.Module<Types.Tensor, Types.Tensor>
	local Wx = Linear(inF, hidF)
	local Wh = Linear(hidF, hidF)
	local self
	self = {
		forward = function(_, x: Types.Tensor, h: Types.Tensor): Types.Tensor
			-- gates
			task.desynchronize()
			local gx = Wx:forward(x) -- {H,B}
			local gh = Wh:forward(h) -- {H,B}
			local B = gx._shape[2]
			local y = Tensor.zeros({ hidF, B }, gx._dtype, (gx._requiresGrad or gh._requiresGrad) == true)
			for i = 1, hidF do
				for b = 1, B do
					local idx = (i - 1) * B + b
					local pre = gx._storage[idx] + gh._storage[idx]
					y._storage[idx] = math.tanh(pre)
				end
			end
			task.synchronize()
			return Tape.makeNode(y, { gx, gh }, function(gy: Types.Tensor)
				local B = y._shape[2]
				if gx._requiresGrad or gh._requiresGrad then
					task.desynchronize()
					if gx._requiresGrad then
						gx._grad = gx._grad or Tensor.zeros(gx._shape, gx._dtype, false)
					end
					if gh._requiresGrad then
						gh._grad = gh._grad or Tensor.zeros(gh._shape, gh._dtype, false)
					end
					for i = 1, hidF do
						for b = 1, B do
							local idx = (i - 1) * B + b
							local dpre = (1 - y._storage[idx] * y._storage[idx]) * gy._storage[idx]
							if gx._requiresGrad then
								gx._grad._storage[idx] += dpre
							end
							if gh._requiresGrad then
								gh._grad._storage[idx] += dpre
							end
						end
					end
					task.synchronize()
				end
			end)
		end,
		parameters = function(_)
			local p = Wx:parameters()
			for _, t in Wh:parameters() do
				table.insert(p, t)
			end
			return p
		end,
	}
	return self
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")
local Util = require("../Util")
local assert = Util.Assert

--@param f Types.Module<Types.Tensor, Types.Tensor>
return function(f: Types.Module<Types.Tensor, Types.Tensor>): Types.Module<Types.Tensor, Types.Tensor>
	return {
		forward = function(_, x: Types.Tensor): Types.Tensor
			local y = f:forward(x)
			assert(y._shape[1] == x._shape[1] and y._shape[2] == x._shape[2], "Residual: shape mismatch")

			local requiresGrad = (y._requiresGrad or x._requiresGrad) == true

			task.desynchronize()
			local out = Tensor.zeros(y._shape, y._dtype, requiresGrad)
			for i = 1, #y._storage do
				out._storage[i] = y._storage[i] + x._storage[i]
			end
			task.synchronize()

			if not requiresGrad then
				return out
			end

			return Tape.makeNode(out, { y, x }, function(gy: Types.Tensor)
				if y._requiresGrad then
					y._grad = y._grad or Tensor.zeros(y._shape)
					task.desynchronize()
					for i = 1, #gy._storage do
						y._grad._storage[i] += gy._storage[i]
					end
					task.synchronize()
				end
				if x._requiresGrad then
					x._grad = x._grad or Tensor.zeros(x._shape)
					task.desynchronize()
					for i = 1, #gy._storage do
						x._grad._storage[i] += gy._storage[i]
					end
					task.synchronize()
				end
			end)
		end,
		parameters = function(self)
			return f:parameters()
		end,
	}
end


--!native
--!optimize 2
local Types = require("../Types")
return function(mods: { Types.Module<Types.Tensor, Types.Tensor> }): Types.Module<Types.Tensor, Types.Tensor>
	return {
		_layers = mods,
		forward = function(_, x: Types.Tensor): Types.Tensor
			local v = x
			for _, m in ipairs(mods) do
				if type(m) == "table" and m.forward then
					v = m:forward(v)
				elseif type(m) == "function" then
					v = m(nil, v)
				else
					error("Sequential: unsupported module entry")
				end
			end
			return v
		end,
		parameters = function(_): { Types.Tensor }
			local out = {}
			for _, m in ipairs(mods) do
				if type(m) == "table" and m.parameters then
					for _, p in ipairs(m:parameters()) do
						table.insert(out, p)
					end
				end
			end
			return out
		end,
	}
end


--!native
--!optimize 2
local SoftmaxOps = require("../ops/Softmax")
local Types = require("../Types")

local S = {}

function S.forward(logits: Types.Tensor): Types.Tensor
	return SoftmaxOps.forward(logits)
end

return S


--!native
--!optimize 2
local Activations = require("./Activations")
local Types = require("../Types")

return function(): Types.Module<Types.Tensor, Types.Tensor>
	return {
		--@param x Types.Tensor @ shape {2H, B}
		--@return Types.Tensor @ shape {H, B}
		forward = function(_, x: Types.Tensor): Types.Tensor
			return Activations.SwiGLUSplit(x)
		end,
		--@return {}
		parameters = function(_): {}
			return {}
		end,
	}
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")

return function(X: Types.Tensor, KH: number, KW: number, stride: number): Types.Tensor
	local C, H, W, B = X._shape[1], X._shape[2], X._shape[3], X._shape[4]
	local Ho = math.floor((H - KH) / stride) + 1
	local Wo = math.floor((W - KW) / stride) + 1
	local Y = Tensor.zeros({ C, Ho, Wo, B }, X._dtype, X._requiresGrad)
	task.desynchronize()
	local inputStorage = X._storage
	local outputStorage = Y._storage
	local sizeScale = 1 / (KH * KW)
	local HW_B = H * W * B
	local W_B = W * B
	for c = 1, C do
		local c_offset = (c - 1) * HW_B
		local c_out_offset = (c - 1) * Ho * Wo * B
		for b = 1, B do
			local b_offset = b - 1
			for i = 1, Ho do
				local i_start = (i - 1) * stride
				local i_out_offset = (i - 1) * Wo * B
				for j = 1, Wo do
					local j_start = (j - 1) * stride
					local sum = 0
					for ki = 0, KH - 1 do
						local row_offset = (i_start + ki) * W_B
						for kj = 0, KW - 1 do
							local col_val = j_start + kj
							-- xi = c_offset + row_offset + col_val * B + b_offset
							local xi = c_offset + row_offset + col_val * B + b_offset
							sum += inputStorage[xi + 1]
						end
					end
					-- yi = c_out_offset + i_out_offset + (j-1)*B + (b-1) + 1
					local yi = c_out_offset + i_out_offset + (j - 1) * B + b_offset + 1
					outputStorage[yi] = sum * sizeScale
				end
			end
		end
	end
	task.synchronize()
	return Tape.makeNode(Y, { X }, function(gy: Types.Tensor)
		if X._requiresGrad then
			X._grad = X._grad or Tensor.zeros(X._shape, X._dtype, false)
			task.desynchronize()
			local scale = 1 / (KH * KW)
			local xGrad = X._grad :: Types.Tensor
			local xGradStorage = xGrad._storage
			local gyStorage = gy._storage
			local HW_B = H * W * B
			local W_B = W * B
			local HoWoB = Ho * Wo * B
			local WoB = Wo * B
			for c = 1, C do
				local c_offset = (c - 1) * HW_B
				local c_out_offset = (c - 1) * HoWoB
				for b = 1, B do
					local b_offset = b - 1
					for i = 1, Ho do
						local i_start = (i - 1) * stride
						local i_out_offset = (i - 1) * WoB
						for j = 1, Wo do
							local yi = c_out_offset + i_out_offset + (j - 1) * B + b_offset + 1
							local gyv = gyStorage[yi] * scale
							local j_start = (j - 1) * stride
							for ki = 0, KH - 1 do
								local row_offset = (i_start + ki) * W_B
								for kj = 0, KW - 1 do
									local col_val = j_start + kj
									local xi = c_offset + row_offset + col_val * B + b_offset
									xGradStorage[xi + 1] += gyv
								end
							end
						end
					end
				end
			end
			task.synchronize()
		end
	end)
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Types = require("../Types")
local Util = require("../Util")
local assert = Util.Assert

-- Matrix multiply: A{M,K} x B{K,N} -> C{M,N}
--@param A Tensor
--@param B Tensor
--@return Tensor
local function matmul(A: Types.Tensor, B: Types.Tensor): Types.Tensor
	local M, K = A._shape[1], A._shape[2]
	local K2, N = B._shape[1], B._shape[2]
	assert(K == K2, "matmul: inner dims mismatch")
	task.desynchronize()
	local C = Tensor.zeros({ M, N }, A._dtype, A._requiresGrad or B._requiresGrad)
	local as, bs, cs = A._storage, B._storage, C._storage
	for i = 1, M do
		local offsetA = (i - 1) * K
		local offsetC = (i - 1) * N
		for k = 1, K do
			local a_ik = as[offsetA + k]
			local offsetB = (k - 1) * N
			for j = 1, N do
				cs[offsetC + j] += a_ik * bs[offsetB + j]
			end
		end
	end
	task.synchronize()
	return C
end

return { matmul = matmul }


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")
local BLAS = require("../ops/BLAS")
local Util = require("../Util")
local assert = Util.Assert

local function copyShape(shape: { number }): { number }
	local out = table.create(#shape)
	for i = 1, #shape do
		out[i] = shape[i]
	end
	return out
end

local function makeMatrixView(tensor: Types.Tensor, rows: number, cols: number): Types.Tensor
	return {
		_storage = tensor._storage,
		_shape = { rows, cols },
		_dtype = tensor._dtype,
		_requiresGrad = tensor._requiresGrad,
		_parents = nil,
		_backprop = nil,
		_grad = nil,
		_size = nil,
	}
end

local function reshapeInPlace(tensor: Types.Tensor, shape: { number }): Types.Tensor
	tensor._shape = copyShape(shape)
	return tensor
end

local function transposeMatrix(matrix: Types.Tensor): Types.Tensor
	local rows, cols = matrix._shape[1], matrix._shape[2]
	local out = Tensor.zeros({ cols, rows }, matrix._dtype, matrix._requiresGrad)
	task.desynchronize()
	local os, ms = out._storage, matrix._storage
	for r = 1, rows do
		local rowOffset = (r - 1) * cols
		for c = 1, cols do
			os[(c - 1) * rows + r] = ms[rowOffset + c]
		end
	end
	task.synchronize()
	return out
end

local function im2col(
	X: Types.Tensor,
	C: number,
	H: number,
	Wd: number,
	B: number,
	KH: number,
	KW: number,
	Ho: number,
	Wo: number
)
	local patch = C * KH * KW
	local cols = Tensor.zeros({ patch, Ho * Wo * B }, X._dtype, false)
	local xs = X._storage
	local cs = cols._storage
	local HWd_B = H * Wd * B
	local Wd_B = Wd * B
	task.desynchronize()
	local col = 0
	for i = 1, Ho do
		for j = 1, Wo do
			for b = 1, B do
				local base = col * patch
				col += 1
				local offset = 0
				local b_minus_1 = b - 1
				for c = 1, C do
					local c_offset = (c - 1) * HWd_B
					for ki = 0, KH - 1 do
						local row_part = (i + ki - 1) * Wd_B -- (i - 1 + ki) * Wd * B
						for kj = 0, KW - 1 do
							local col_part = (j + kj - 1) * B -- (j - 1 + kj) * B
							local xi = c_offset + row_part + col_part + b_minus_1
							cs[base + offset + 1] = xs[xi + 1]
							offset += 1
						end
					end
				end
			end
		end
	end
	task.synchronize()
	return cols
end

local function col2im(
	cols: Types.Tensor,
	C: number,
	H: number,
	Wd: number,
	B: number,
	KH: number,
	KW: number,
	Ho: number,
	Wo: number,
	dtype: Types.DType
): Types.Tensor
	local patch = C * KH * KW
	local out = Tensor.zeros({ C, H, Wd, B }, dtype, false)
	local xs = out._storage
	local cs = cols._storage
	local HWd_B = H * Wd * B
	local Wd_B = Wd * B
	task.desynchronize()
	local col = 0
	for i = 1, Ho do
		for j = 1, Wo do
			for b = 1, B do
				local base = col * patch
				col += 1
				local offset = 0
				local b_minus_1 = b - 1
				for c = 1, C do
					local c_offset = (c - 1) * HWd_B
					for ki = 0, KH - 1 do
						local row_part = (i + ki - 1) * Wd_B
						for kj = 0, KW - 1 do
							local col_part = (j + kj - 1) * B
							local xi = c_offset + row_part + col_part + b_minus_1
							xs[xi + 1] += cs[base + offset + 1]
							offset += 1
						end
					end
				end
			end
		end
	end
	task.synchronize()
	return out
end

local function addInto(dst: Types.Tensor, src: Types.Tensor)
	task.desynchronize()
	local ds, ss = dst._storage, src._storage
	for i = 1, #ds do
		ds[i] += ss[i]
	end
	task.synchronize()
end

--@param X Tensor
--@param W Tensor -- {F,C,KH,KW}
--@return Tensor -- {F,Ho,Wo,B}
return function(X: Types.Tensor, W: Types.Tensor): Types.Tensor
	local C, H, Wd, B = X._shape[1], X._shape[2], X._shape[3], X._shape[4]
	local F, Cw, KH, KW = W._shape[1], W._shape[2], W._shape[3], W._shape[4]
	assert(C == Cw, "conv2d: channel mismatch")
	local Ho = H - KH + 1
	local Wo = Wd - KW + 1
	assert(Ho > 0 and Wo > 0, "conv2d: kernel larger than input")
	task.desynchronize()
	local patch = C * KH * KW
	local columns = Ho * Wo * B
	local cols = im2col(X, C, H, Wd, B, KH, KW, Ho, Wo)
	local weightView = makeMatrixView(W, F, patch)
	local yMat = BLAS.matmul(weightView, cols)
	local Y = reshapeInPlace(yMat, { F, Ho, Wo, B })
	task.synchronize()
	return Tape.makeNode(Y, { X, W }, function(gy: Types.Tensor)
		local gyMat = makeMatrixView(gy, F, columns)
		if W._requiresGrad then
			if not W._grad then
				W._grad = Tensor.zeros({ F, C, KH, KW }, W._dtype, false)
			end
			local colsT = transposeMatrix(cols)
			local dW = BLAS.matmul(gyMat, colsT)
			local gradW = W._grad :: Types.Tensor
			local gradView = makeMatrixView(gradW, F, patch)
			addInto(gradView, dW)
		end
		if X._requiresGrad then
			if not X._grad then
				X._grad = Tensor.zeros({ C, H, Wd, B }, X._dtype, false)
			end
			local weightT = transposeMatrix(weightView)
			local dXcols = BLAS.matmul(weightT, gyMat)
			local dX = col2im(dXcols, C, H, Wd, B, KH, KW, Ho, Wo, X._dtype)
			local gradX = X._grad :: Types.Tensor
			addInto(gradX, dX)
		end
	end)
end