/* SPDX-License-Identifier: AGPL-3.0-or-later AND Apache-2.0 */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * Portions derived from wkv7_cuda_fp32.cu in RWKV-LM
 * (https://github.com/BlinkDL/RWKV-LM), by Bo Peng (BlinkDL),
 * licensed under the Apache License, Version 2.0. Modified by
 * David (Corey) Herrera, 2026. See NOTICE and LICENSES/Apache-2.0.txt.
 */
/*
 * rx_pretrain.h — RWKV-7 pretraining kernels (backward-through-time)
 *
 * Prefix: rx_ (RWKV-X).  This is the start of the pretrain module; for now it
 * contains ONLY the WKV7 backward-through-time (the counterpart to the forward
 * `rx_wkv7_one` in rwkvx.c) plus a full-sequence forward used to drive it and
 * the gradient checks.  The CE loop / base-unfreeze / full pretrain loop come
 * later and are intentionally NOT here.
 *
 * WKV7 inner recurrence (per head, per timestep; identical to rx_wkv7_one):
 *     sa[i]   = sum_j  a[j] * S_prev[i][j]           (a = nkk = -kk)
 *     S[i][j] = S_prev[i][j]*w[j] + v[i]*k[j] + sa[i]*b[j]   (b = kka = kk*iclr)
 *     y[i]    = sum_j  S[i][j] * r[j]
 * State element convention matches foundry `s_d[(h*N+i)*N + j]`: row i is the
 * value/output dim, col j is the key dim.
 *
 * Backward equations transcribed from BlinkDL's Apache-2.0 CUDA
 * wkv7_cuda_fp32.cu backward_kernel via the verified fp64 PyTorch oracle
 * (Ternary_scale_1-10m/reference/wkv7/wkv7_reference.py).
 *
 * IMPORTANT chain-rule hand-offs (see wkv7_backward_reference.md §5):
 *   1. `dw` here is the grad w.r.t. the EFFECTIVE decay w_vec (already in (0,1)),
 *      NOT the log-space pre-activation.  We emit bare dw; the surrounding
 *      time-mix backward chains it through exp(-0.606531*sigmoid(...)).  We do
 *      NOT apply the CUDA kernel's `dw*wi*wi_fac` factor.
 *   2. `da`/`db` are grads w.r.t. the composites nkk(=a) and kka(=b).  The
 *      recombination onto shared kk and a_vec (and a_vec's second path into
 *      k_mod) is the time-mix backward's job, out of scope here.
 *   3. `dS` after the reverse scan is the grad w.r.t. the initial state carry
 *      S_{-1}.  Zero it for full-sequence-from-zero; it is the true exerciser
 *      of the reverse scan and is checked via finite differences.
 */
#ifndef RWKVX_PRETRAIN_H
#define RWKVX_PRETRAIN_H

#include <stdint.h>

#ifdef __cplusplus
extern "C" {
#endif

/*
 * Full-sequence WKV7 forward.  Layout for the per-timestep vectors is
 * [B][T][H][N] flattened row-major: index = (((b*T + t)*H + h)*N + n).
 * State layout is [B][H][N][N]: index = (((b*H + h)*N + i)*N + j).
 *
 *   s_init : initial state carry [B*H*N*N], or NULL for a zero start.
 *   y_out  : forward output [B*T*H*N], or NULL if not needed.
 *   s_final: final state after T steps [B*H*N*N], or NULL if not needed.
 *
 * Returns 0 on success, -1 on allocation failure.
 */
int rx_wkv7_forward_seq(
    const float *r, const float *w, const float *k, const float *v,
    const float *a, const float *b, const float *s_init,
    float *y_out, float *s_final,
    int B, int T, int n_head, int head_size);

/*
 * WKV7 backward-through-time (reverse scan).  Same layouts as above.
 *
 *   dy     : upstream grad dL/dy [B*T*H*N].
 *   s_init : initial state carry used in the forward [B*H*N*N], or NULL (zero).
 *   dr..db : output grads [B*T*H*N] each (dw is w.r.t. effective decay, see §5).
 *   dS_out : output grad w.r.t. the initial state carry S_{-1} [B*H*N*N], or
 *            NULL if training full-sequence-from-zero (grad discarded).
 *
 * The forward is recomputed internally and every pre-step state S_prev is
 * cached (correctness-first; the CUDA kernel reconstructs to save memory).
 * Returns 0 on success, -1 on allocation failure.
 */
int rx_wkv7_backward(
    const float *r, const float *w, const float *k, const float *v,
    const float *a, const float *b, const float *dy, const float *s_init,
    float *dr, float *dw, float *dk, float *dv, float *da, float *db,
    float *dS_out,
    int B, int T, int n_head, int head_size);

/* ================================================================
 * RWKV-7 TIME-MIX BLOCK backward (wraps the WKV7 backward above).
 *
 * This is the counterpart of rwkvx.c `rwkv7_time_mix_fast`: it covers the
 * WHOLE time-mix block (token-shift/lerp, the R/K/V/output projections, the
 * low-rank w/a/g/v LoRA gates, key L2-norm + key-modification, value-residual
 * gating, the WKV7 recurrence, group-norm, the r*k*r_k bonus and the g gate),
 * EXCLUDING the pre-block ln1 and the residual add.  The block input is the
 * (already LN'd) hidden state; the output feeds the residual.
 *
 * It computes dL/d(block input) and grads for every learnable time-mix param,
 * internally recomputing the fp32 forward (caching activations) and delegating
 * the recurrence to rx_wkv7_backward.  The two §5 hand-offs live here:
 *   - `dw` (effective decay) is chained through exp(-0.606531*sigmoid(...)).
 *   - `da`/`db` (grads w.r.t. nkk=-kk and kka=kk*a_vec) are recombined onto the
 *     shared kk and a_vec (which also feeds k_mod).
 *
 * Layouts: block input x, output, dy, dx are [B][T][C] row-major (C=n_head*
 * head_size).  Projection weights are row-major [in,out] exactly as
 * rwkvx.c's fast_matvec consumes them.  v_first is [B][T][C] (per-token; for
 * layer_id==0 it is an OUTPUT copy of v and its grad is ignored).
 */
typedef struct {
    int c;                  /* n_head*head_size */
    int n_head, head_size;
    int Dw, Da, Dg, Dv;     /* LoRA inner dims (Dv ignored when layer_id==0) */
    int layer_id;
    const float *x_r, *x_k, *x_v, *x_w, *x_a, *x_g;         /* [c] lerps */
    const float *receptance_w, *key_w, *value_w, *output_w; /* [c*c] [in,out] */
    const float *w0, *w1, *w2;   /* w0[c], w1[c*Dw], w2[Dw*c] */
    const float *a0, *a1, *a2;   /* a0[c], a1[c*Da], a2[Da*c] */
    const float *g1, *g2;        /* g1[c*Dg], g2[Dg*c] */
    const float *v0, *v1, *v2;   /* v0[c], v1[c*Dv], v2[Dv*c] (layer_id>0) */
    const float *k_k, *k_a;      /* [c] */
    const float *r_k;            /* [c] */
    const float *ln_x_w, *ln_x_b;/* [c] group-norm affine */
} rx_tm_params;

/* Grad accumulators — same shapes as rx_tm_params, non-const.  Every pointer
 * must be caller-allocated; rx_timemix_backward ZEROES them at entry.  Pointers
 * for params absent in the configuration (v0/v1/v2 at layer 0) may be NULL. */
typedef struct {
    float *x_r, *x_k, *x_v, *x_w, *x_a, *x_g;
    float *receptance_w, *key_w, *value_w, *output_w;
    float *w0, *w1, *w2;
    float *a0, *a1, *a2;
    float *g1, *g2;
    float *v0, *v1, *v2;
    float *k_k, *k_a;
    float *r_k;
    float *ln_x_w, *ln_x_b;
} rx_tm_grads;

/* fp32 time-mix forward over a full sequence (per-batch, T tokens).
 *   x           : block input [B*T*c] (LN'd hidden state)
 *   x_prev_init : token-shift carry at t=0 [B*c], or NULL for zero.
 *   v_first     : [B*T*c].  layer_id==0: WRITTEN (copy of v).  layer_id>0: READ.
 *   out         : block output [B*T*c] (pre-residual).
 * Returns 0 on success, -1 on allocation failure. */
int rx_timemix_forward_seq(
    const rx_tm_params *P, const float *x, const float *x_prev_init,
    float *v_first, float *out, int B, int T);

/* Time-mix backward.  dy is dL/d(out) [B*T*c].  Produces dx [B*T*c], all param
 * grads in G, and (optional) d_xprev_init [B*c] / d_v_first [B*T*c].  NULLs for
 * unused outputs.  Returns 0 on success, -1 on allocation failure. */
int rx_timemix_backward(
    const rx_tm_params *P, const float *x, const float *x_prev_init,
    const float *v_first, const float *dy,
    rx_tm_grads *G, float *dx, float *d_xprev_init, float *d_v_first,
    int B, int T);

/* ================================================================
 * RWKV-7 CHANNEL-MIX (ReLU^2 FFN) backward.
 *
 * Counterpart of rwkvx.c `rwkv7_channel_mix_fast`: token-shift/lerp, the key
 * projection [C,ffn_dim], the squared-ReLU nonlinearity (k = relu(fk)^2), and
 * the value projection [ffn_dim,C].  EXCLUDING the pre-block ln2 and the
 * residual add (those live in the block).  Input x is the (already LN'd)
 * hidden state [B][T][c]; output feeds the residual.
 *
 * Layouts match the time-mix block: x/out/dy/dx are [B][T][c] row-major.
 * ffn_key_w is row-major [c,ffn_dim]; ffn_value_w is [ffn_dim,c] (both as
 * fast_matvec consumes them).  The token-shift carry (ffn_x_prev) at t=0 is
 * x_prev_init [B][c] or NULL for a zero start; its grad is d_xprev_init.
 */
typedef struct {
    int c;
    int ffn_dim;
    const float *ffn_x_k;      /* [c] token-shift lerp */
    const float *ffn_key_w;    /* [c*ffn_dim] [in,out] */
    const float *ffn_value_w;  /* [ffn_dim*c] [in,out] */
} rx_cm_params;

typedef struct {
    float *ffn_x_k;
    float *ffn_key_w;
    float *ffn_value_w;
} rx_cm_grads;

/* fp32 channel-mix forward over a full sequence.  Returns 0 / -1(alloc). */
int rx_channelmix_forward_seq(
    const rx_cm_params *P, const float *x, const float *x_prev_init,
    float *out, int B, int T);

/* Channel-mix backward.  dy is dL/d(out) [B*T*c].  Produces dx [B*T*c], all
 * param grads in G (ZEROED at entry), and optional d_xprev_init [B*c].
 * Returns 0 / -1(alloc). */
int rx_channelmix_backward(
    const rx_cm_params *P, const float *x, const float *x_prev_init,
    const float *dy, rx_cm_grads *G, float *dx, float *d_xprev_init,
    int B, int T);

/* ================================================================
 * LayerNorm backward (the block's ln1/ln2; foundry `fast_layer_norm`,
 * eps = 1e-5, normalised over the full C channel).  Standard LN backward.
 *
 * ⚠ The mean/variance reductions cancel in fp32 exactly as the L2-norm did;
 * they are accumulated in DOUBLE locally so gamma/beta/dx stay clean.
 *
 * x/dy/dx are [B][T][C] row-major.  dw/db are [C] and are ZEROED at entry
 * then accumulated across all B*T tokens.  dx is written (=), not accumulated.
 * Returns 0 / -1(alloc).
 * ================================================================ */
int rx_layernorm_forward_seq(
    const float *x, const float *w, const float *b, float *out,
    int B, int T, int C);

int rx_layernorm_backward(
    const float *x, const float *w, const float *b, const float *dy,
    float *dx, float *dw, float *db, int B, int T, int C);

/* ================================================================
 * RWKV-7 BLOCK forward/backward — the full residual block.
 * Mirrors rwkvx.c `rwkv7_block_forward_fast`:
 *   x -> ln1 -> time-mix -> +x (residual) -> ln2 -> channel-mix -> +x1
 * Backward walks it in reverse, threading residual grads to BOTH the block
 * path and the skip.  Produces dL/d(block input) + every block param grad
 * (time-mix + channel-mix + ln1 + ln2).
 * ================================================================ */
typedef struct {
    rx_tm_params tm;
    rx_cm_params cm;
    const float *ln1_w, *ln1_b;   /* [c] */
    const float *ln2_w, *ln2_b;   /* [c] */
} rx_block_params;

typedef struct {
    rx_tm_grads tm;
    rx_cm_grads cm;
    float *ln1_w, *ln1_b, *ln2_w, *ln2_b;
} rx_block_grads;

/* Block forward.  v_first [B*T*c]: WRITTEN at layer 0 (copy of v), READ at
 * layer>0.  tm_xprev_init / cm_xprev_init are the two token-shift carries
 * [B*c] (over ln1-out and ln2-out respectively) or NULL for zero.  Returns
 * 0 / -1(alloc). */
int rx_block_forward_seq(
    const rx_block_params *P, const float *x,
    const float *tm_xprev_init, const float *cm_xprev_init,
    float *v_first, float *out, int B, int T);

/* Block backward.  dy is dL/d(out) [B*T*c].  Produces dx [B*T*c], all grads in
 * G (sub-grads ZEROED by their sub-backwards; ln grads zeroed here), and the
 * optional carry / v_first grads.  Returns 0 / -1(alloc).
 *
 * d_v_first is dual-purpose (see rx_timemix_backward): for a layer>0 block it is
 * an OUTPUT (grad w.r.t. the v_first input, zeroed then written); for the
 * layer-0 block (which PRODUCES v_first) it is an optional INPUT carrying the
 * accumulated downstream grad on v_first, injected into the value path.  Pass
 * NULL when there is no v_first coupling (single RWKV layer / legacy). */
int rx_block_backward(
    const rx_block_params *P, const float *x,
    const float *tm_xprev_init, const float *cm_xprev_init,
    const float *v_first, const float *dy,
    rx_block_grads *G, float *dx,
    float *d_tm_xprev_init, float *d_cm_xprev_init, float *d_v_first,
    int B, int T);

/* ================================================================
 * SPARSE (MoBA-style) ATTENTION BLOCK forward/backward — BASE weights.
 *
 * Counterpart of rwkvx.c `attn_block_forward_fast` for the full residual
 * attention block (RWKV-x interleaves these with the RWKV-7 blocks above):
 *   h1  = ln1(x)
 *   q,k,v = h1 @ {Wq,Wk,Wv}                    (base projections, [c,c])
 *   att = MoBA-sparse causal-softmax attention(q,k,v)   scale = 1/sqrt(head_size)
 *   att_out = att @ Wo ;  x1 = x + att_out       (residual)
 *   h2  = ln2(x1)
 *   cm  = channel-mix(h2) ;  out = x1 + cm        (residual)
 *
 * This computes dL/d(block input) and grads for EVERY learnable attention-block
 * base parameter (Q/K/V/O projections, ln1/ln2 affine, and the channel-mix
 * params) — the full base, not just the LoRA of rx_train.c.  The softmax scale
 * is the fixed constant 1/sqrt(head_size) (not learnable); the forward has no
 * RoPE, so there is nothing to invert there.
 *
 * MoBA top-k CHUNK SELECTION IS HARD / NON-DIFFERENTIABLE.  The selected chunk
 * set is treated as FIXED and the backward runs the standard sparse-attention
 * gradient over exactly the selected keys (keys of unselected chunks get
 * attention weight 0, so they drop out of every reduction automatically).  For
 * a finite-difference check, perturbations must not flip the selection — test
 * in a regime where it is stable (e.g. moba_top_k >= #candidate chunks so all
 * are selected => plain causal attention).  See tests/test_attn_bwd.c.
 *
 * Selection rule (per query t, per head): keys are partitioned into chunks of
 * `chunk_size`.  The chunk containing t is always attended (causally, keys
 * s<=t).  Each earlier complete chunk gets a gate = <q_t, mean_key(chunk)>; the
 * top `moba_top_k` earlier chunks are selected.  moba_top_k <= 0 or
 * >= #candidate chunks both mean "keep all" (full causal).
 * ================================================================ */
typedef struct {
    int c;                     /* n_head*head_size */
    int n_head, head_size;
    int chunk_size;            /* MoBA chunk size (>=1) */
    int moba_top_k;            /* # earlier chunks kept; <=0 => keep all */
    const float *ln1_w, *ln1_b;      /* [c] pre-attention LN affine */
    const float *receptance_w;       /* [c*c] Q proj [in,out] */
    const float *key_w;              /* [c*c] K proj */
    const float *value_w;            /* [c*c] V proj */
    const float *output_w;           /* [c*c] O proj */
    const float *ln2_w, *ln2_b;      /* [c] pre-FFN LN affine */
    rx_cm_params cm;                 /* channel-mix (ffn_x_k, ffn_key_w, ffn_value_w) */
} rx_attn_params;

typedef struct {
    float *ln1_w, *ln1_b;
    float *receptance_w, *key_w, *value_w, *output_w;
    float *ln2_w, *ln2_b;
    rx_cm_grads cm;
} rx_attn_grads;

/* fp32 forward over a full sequence (per batch, causal, MoBA-sparse).
 *   x              : block input [B*T*c]
 *   cm_xprev_init  : channel-mix token-shift carry at t=0 [B*c] or NULL(zero)
 *   out            : block output [B*T*c]
 * Returns 0 / -1(alloc). */
int rx_attn_forward_seq(
    const rx_attn_params *P, const float *x, const float *cm_xprev_init,
    float *out, int B, int T);

/* Attention-block backward with BASE-weight gradients.  dy is dL/d(out)
 * [B*T*c].  Produces dx [B*T*c] and every param grad in G (all ZEROED at
 * entry), plus optional d_cm_xprev_init [B*c].  The forward is recomputed
 * internally; the MoBA selection is recomputed from the base weights and held
 * fixed through the attention backward.  Returns 0 / -1(alloc). */
int rx_attn_backward_full(
    const rx_attn_params *P, const float *x, const float *cm_xprev_init,
    const float *dy, rx_attn_grads *G, float *dx, float *d_cm_xprev_init,
    int B, int T);

/* Exported test hooks: the ISOLATED MoBA attention core (scores->softmax->AV),
 * mirrored on GPU in cuda_kernels/moba_attn.cu and parity-checked by
 * cuda_kernels/tests/test_moba_attn_cuda.cu.  Inputs are the already-projected
 * per-token q,k,v [B*T*c]; forward fills softmax weights p [B*n_head*T*T] and
 * attention output y [B*T*c]; backward recomputes the MoBA selection + p from
 * q,k,v and accumulates dq,dk,dv [B*T*c] (all zeroed at entry) from dy=d(y). */
int rx_attn_score_av_forward(const rx_attn_params *P, const float *q, const float *k,
    const float *v, int B, int T, float *p, float *y);
int rx_attn_score_av_backward(const rx_attn_params *P, const float *q, const float *k,
    const float *v, const float *dy, int B, int T,
    float *dq, float *dk, float *dv);

/* ================================================================
 * FULL RWKV-x MODEL: embedding, head/unembedding, cross-entropy, from-scratch
 * init, and the whole-model forward/backward assembly.  Single sequence
 * (B=1); the underlying block kernels are called per-sequence.
 *
 * Architecture mirrors rwkvx.c rx_forward:
 *   x0 = emb[token]                       (raw embedding; foundry fuses ln0 into
 *                                          emb at load, so a from-scratch random
 *                                          emb absorbs it — no separate ln0 here)
 *   for each block (RWKV-7 and sparse-attention, interleaved per the RWKV-x
 *       pattern — interval = n_rwkv/n_attn, attention after each RWKV group):
 *       x = block(x)                      (residual stream)
 *   h  = ln_out(x)                        (final LayerNorm)
 *   logits = h @ head_w                   ([C,vocab]; head is NOT weight-tied to
 *                                          emb — foundry stores rwkv.head.weight
 *                                          separately)
 *   loss = mean_t CE(logits[t], target[t])
 *
 * v_first coupling: the RWKV-7 block at rwkv-array index 0 produces v_first;
 * later RWKV-7 blocks consume it.  The reverse pass accumulates the grad on
 * v_first from the consumers and injects it into producer block 0 via the
 * rx_timemix_backward d_v_first-as-input path (see above).
 * ================================================================ */

/* ---- Embedding (gather / scatter-add) ---- */
void rx_embedding_forward(const float *emb, const int *tokens,
                          int T, int C, float *out);           /* out[T*C] */
void rx_embedding_backward(const int *tokens, int T, int C, int vocab,
                           const float *d_out, float *d_emb);  /* d_emb ZEROED then scatter-add */

/* ---- Head / unembedding.  head is [C,vocab] row-major (as fast_matvec eats). ---- */
void rx_head_forward(const float *x, const float *head,
                     int T, int C, int vocab, float *logits);  /* logits[T*vocab] */
void rx_head_backward(const float *x, const float *head, const float *dlogits,
                      int T, int C, int vocab,
                      float *dx, float *d_head);                /* dx written(=); d_head ZEROED then accumulate */

/* ---- Fused log-softmax cross-entropy over the vocab.  Mean over valid tokens
 * (targets != ignore_index; pass ignore_index<0 to disable).  log-sum-exp is
 * double-accumulated for numerical cleanliness. ---- */
double rx_cross_entropy_forward(const float *logits, const int *targets,
                                int T, int vocab, int ignore_index);
/* Returns loss and writes dlogits = (softmax - onehot)/n_valid (ignored token
 * rows are zero). */
double rx_cross_entropy_backward(const float *logits, const int *targets,
                                 int T, int vocab, int ignore_index,
                                 float *dlogits);               /* dlogits[T*vocab] */

/* ---- From-scratch model ---- */
typedef struct {
    int n_embd, n_head, head_size;   /* head_size*n_head == n_embd */
    int n_rwkv_layer, n_attn_layer;
    int vocab_size;
    int Dw, Da, Dg, Dv;              /* time-mix LoRA inner dims */
    int ffn_dim;                     /* channel-mix hidden dim */
    int chunk_size, moba_top_k;      /* attention MoBA config (>=candidates => full causal) */
} rx_pt_config;

typedef struct rx_ptmodel rx_ptmodel;   /* opaque; owns all params + grads */

/* Deterministic seeded init of ALL params (splitmix64 PRNG).  Init scheme:
 *   linear projections & LoRA matrices : uniform(-0.5,0.5) * 0.3
 *   token-shift / gate lerps, w0/a0/v0/k_k/k_a/r_k : uniform(-0.5,0.5) * 0.5
 *   LayerNorm / group-norm gamma       : 1 + 0.2*uniform(-0.5,0.5)
 *   LayerNorm / group-norm beta        : 0.2*uniform(-0.5,0.5)
 *   embedding                          : uniform(-0.5,0.5) * 0.5
 *   head                               : uniform(-0.5,0.5) * 0.3
 * Small-scale, non-degenerate, keeps CE finite; a production trainer would
 * refine these.  Returns NULL on bad config / alloc failure. */
rx_ptmodel* rx_model_new_random(const rx_pt_config *pc, uint64_t seed);
void        rx_ptmodel_free(rx_ptmodel *m);

/* Introspection for the gradient check / reporting. */
int   rx_ptmodel_num_params(const rx_ptmodel *m);
/* Per-tensor accessor: name, mutable data, matching analytic grad, element count. */
int   rx_ptmodel_param(rx_ptmodel *m, int i, const char **name,
                       float **data, float **grad, int *n);
rx_pt_config rx_ptmodel_config(const rx_ptmodel *m);
/* Direct views onto a block's live (possibly-perturbed) params — for an fp64
 * mirror / gradient check.  idx is the index within the RWKV / attn array. */
const rx_block_params* rx_ptmodel_rwkv(rx_ptmodel *m, int idx);
const rx_attn_params*  rx_ptmodel_attn(rx_ptmodel *m, int idx);

/* ---- Activation cache for backward ---- */
typedef struct {
    int T, C, vocab, n_blocks;
    int   *tokens;      /* [T] */
    float *block_in;    /* [(n_blocks+1)*T*C]: entry residual per block; [n_blocks]=final */
    float *v_first;     /* [T*C] */
    float *lnout;       /* [T*C] ln_out output (head input) */
    float *logits;      /* [T*vocab] */
} rx_pt_acts;

rx_pt_acts* rx_pt_acts_new(const rx_ptmodel *m, int T);
void        rx_pt_acts_free(rx_pt_acts *a);

/* Full forward: gather emb -> interleaved blocks -> ln_out -> head -> logits.
 * Fills acts (incl. logits).  Returns 0 / -1(alloc). */
int    rx_pretrain_forward(rx_ptmodel *m, const int *tokens, int T,
                           rx_pt_acts *acts);

/* Full backward: CE -> head -> ln_out -> reverse blocks (interleaved) ->
 * embedding scatter-add.  Accumulates every param grad into the model's grad
 * buffers (zeroed by the sub-backwards / here) and returns the scalar loss. */
double rx_pretrain_backward(rx_ptmodel *m, const rx_pt_acts *acts,
                            const int *targets, int T, int ignore_index);

/* Zero all analytic grad buffers (call before a backward if reusing). */
void   rx_ptmodel_zero_grad(rx_ptmodel *m);

#ifdef __cplusplus
}
#endif

#endif /* RWKVX_PRETRAIN_H */