tenzr::txpr¶

Namespaces¶

Classes¶

Types¶

Functions¶

Types Documentation¶

enum OpCode¶

using TapeCompiledModule¶

using tenzr::txpr::TapeCompiledModule = typedef TxprCompiledModule;

Compatibility alias for TxprCompiledModule.

using TapeEmitter¶

using tenzr::txpr::TapeEmitter = typedef TxprEmitter;

Compatibility alias for TxprEmitter.

using Instruction¶

using tenzr::txpr::Instruction = typedef TxprEqn;

using LinearTape¶

using tenzr::txpr::LinearTape = typedef Txpr;

Functions Documentation¶

function Cond¶

TENZR_APIArray Cond(
    const Array & pred,
    const std::function< Array()> & true_fn,
    const std::function< Array()> & false_fn
)

Parameters:

• pred Boolean scalar tensor.
• true_fn Branch taken when pred is true.
• false_fn Branch taken when pred is false.

Return: The selected branch's result, with both branches required to agree in shape and dtype.

Conditional branching operator analogous to jax.lax.cond. Evaluates true_fn() if pred is true, else false_fn(). Both branches must return arrays of identical shape and dtype.

Select between two branches based on a boolean predicate.

function Cond¶

TENZR_APIArray Cond(
    const Array & pred,
    const std::function< Array(const Array &)> & true_fn,
    const std::function< Array(const Array &)> & false_fn,
    const Array & operand
)

Operand-carrying overload of cond, passing operand to the selected branch.

Parameters:

• pred Boolean scalar tensor.
• true_fn Branch taken when pred is true.
• false_fn Branch taken when pred is false.
• operand Value passed to the selected branch.

Return: The selected branch's result.

Predicated cond variant that forwards a shared operand to the chosen branch.

function Scan¶

TENZR_APIScanResult Scan(
    const std::function< ScanResult(const Array &carry, const Array &x)> & body_fn,
    const Array & init_carry,
    const Array & xs,
    int64_t length =-1
)

Parameters:

• body_fn Step function mapping (carry, x) to (next_carry, y).
• init_carry Initial carry value.
• xs Input sequence, scanned along axis 0.
• length Number of steps; uses xs.dim(0) when negative.

Return: Final carry and the stacked per-step outputs.

Scan a function over leading array dimensions carrying state, analogous to jax.lax.scan. Given body_fn: (carry, x) -> (new_carry, y), loops over slices of xs along axis 0, carrying and updating state while collecting stacked outputs y. Returns: pair of (final_carry, stacked_ys).

Run a carry-based fold over the leading axis of xs, collecting per-step outputs.

function WhileLoop¶

TENZR_APIArray WhileLoop(
    const std::function< Array(const Array &)> & cond_fn,
    const std::function< Array(const Array &)> & body_fn,
    const Array & init_val
)

Parameters:

• cond_fn Predicate returning a scalar boolean on the current value.
• body_fn Step function producing the next value.
• init_val Initial loop value.

Return: Final loop value; capped at 100,000 iterations.

Dynamic iterative while loop analogous to jax.lax.while_loop. Loops while cond_fn(val) is true, updating val = body_fn(val). Does not require static loop unrolling.

Repeat body_fn from init_val until cond_fn returns false.

function IsZeroDerivativeCustomOp¶

bool IsZeroDerivativeCustomOp(
    std::string_view op
)

Names of custom ops whose value is piecewise-constant (zero gradient / zero tangent).

function HasDerivativeRule¶

bool HasDerivativeRule(
    std::string_view op
)

Whether a custom-op name is in the VJP table.

function HasJvpRule¶

bool HasJvpRule(
    std::string_view op
)

Whether a custom-op name is in the JVP table (a subset of the VJP table).

function CustomOpVjp¶

void CustomOpVjp(
    std::string_view op,
    const std::vector< Array > & inputs,
    const Array & grad_out,
    std::vector< Array > & out
)

Vector-Jacobian product for a custom op: appends the input gradients of inputs to out, unbroadcast to each input's shape, given the output cotangent.

function CustomOpJvp¶

Array CustomOpJvp(
    std::string_view op,
    const std::vector< Array > & primals,
    const std::vector< Array > & tangents
)

Forward-mode tangent for a custom op: returns d(out)/d(inputs) applied to the tangents.

function CustomVjp¶

TENZR_APICustomVJPFunction CustomVjp(
    std::string name,
    std::function< Array(const Array &)> fn,
    std::function< std::pair< Array, Array >(const Array &)> fwd,
    std::function< Array(const Array &, const Array &)> bwd
)

Returns fn wrapped with a custom VJP rule registered under name.

Parameters:

• name Function name.
• fn Eager forward implementation.
• fwd Forward rule used under tracing.
• bwd Backward rule.

Return: The constructed function.

Build a CustomVJPFunction from a name and its forward/backward implementations.

function RegisterCustomVjp¶

TENZR_API void RegisterCustomVjp(
    const std::string & name,
    std::function< Array(const Array &, const Array &)> bwd_fn
)

Registers bwd_fn as the backward rule named name.

Parameters:

• name Operation name.
• bwd_fn Backward rule.

Register only a backward rule for an existing operation.

function SetGradTracingActive¶

TENZR_API void SetGradTracingActive(
    bool active
)

Enables or disables gradient tracing globally.

Parameters:

• active New flag value.

Enable or disable gradient tracing on this thread.

function IsGradTracingActive¶

TENZR_API bool IsGradTracingActive()

Returns whether gradient tracing is currently active.

Return: Current flag value.

Whether gradient tracing is active on this thread.

function Grad¶

std::unordered_map< uint64_t, Array > Grad(
    const Array & loss,
    const std::vector< Array > & targets
)

Parameters:

• loss Scalar loss to differentiate.
• targets Arrays whose gradients are requested.

Return: Map from target id to its gradient.

Tensor-level reverse-mode automatic differentiation: Computes gradients of a scalar loss with respect to a list of target Arrays. Supports both dynamic tracer tape differentiation (VJP) and AST graph traversal.

Compute gradients of a scalar loss with respect to several targets.

function Grad¶

Array Grad(
    const Array & loss,
    const Array & target
)

Convenience overload for computing gradient of a scalar loss with respect to a single target.

function Vjp¶

VjpResult Vjp(
    const std::function< Array(const std::vector< Array > &)> & fn,
    const std::vector< Array > & primals
)

Multi-argument vector-Jacobian product (matching jax.vjp): given f: vector -> scalar Array, returns {y = f(primals), pullback} where pullback(cotangent) returns the cotangents for each primal, in order.

function Grad¶

std::function< Array(const Array &)> Grad(
    const std::function< Array(const Array &)> & fn
)

Type-erased functional gradient transform: Given a function f: Array -> scalar Array, returns df/dx: Array -> Array.

function ValueAndGrad¶

std::function< ValueAndGradResult(const Array &)> ValueAndGrad(
    const std::function< Array(const Array &)> & fn
)

Type-erased functional value-and-grad transform: Given a function f: Array -> scalar Array, returns a function that computes std::pair{f(x), df/dx(x)}.

function Grad¶

template <typename Fn >
inline std::function< Array(const Array &)> Grad(
    Fn && fn
)

Functional gradient transform (matching jax.grad). Constrained via C++20 std::invocable concept and excluding exact std::function to prevent recursion.

function ValueAndGrad¶

template <typename Fn >
inline std::function< ValueAndGradResult(const Array &)> ValueAndGrad(
    Fn && fn
)

Functional value-and-grad transform (matching jax.value_and_grad). Constrained via C++20 std::invocable concept and excluding exact std::function to prevent recursion.

function ExecuteJittedExecutable¶

ExecutionResult ExecuteJittedExecutable(
    const std::shared_ptr< CachedExecutable > & cached,
    const std::vector< Array > & inputs
)

Dispatches an already compiled executable with concrete input Array device buffers.

Parameters:

• cached Compiled executable to run.
• inputs Input arrays, uploaded lazily and zero-copy when already resident.

Return: Device-resident execution result.

Upload inputs to the device and execute a cached PJRT executable.

function CompileTxprModule¶

std::shared_ptr< CachedExecutable > CompileTxprModule(
    uint64_t hash,
    const Txpr & txpr
)

Lowers a Txpr to StableHLO MLIR, compiles via PJRT, and inserts into JitCache.

Parameters:

• hash Structural hash identifying the graph.
• txpr Traced IR graph to compile.

Return: The compiled, cached executable.

Emit the Txpr as StableHLO, compile it through PJRT, and cache the executable.

function CompileTapeModule¶

inline std::shared_ptr< CachedExecutable > CompileTapeModule(
    uint64_t hash,
    const Txpr & txpr
)

Compatibility alias for compile_txpr_module.

function Jit¶

template <typename Fn >
auto Jit(
    Fn && fn
)

Function decorator analogous to jax.jit.

function Jvp¶

TENZR_APIJvpResult Jvp(
    const std::function< Array(const Array &)> & fn,
    const Array & primal,
    const Array & tangent
)

Parameters:

• fn Function to differentiate.
• primal Point at which to evaluate.
• tangent Direction to push forward.

Return: The primal output and its JVP.

Computes a forward-mode Jacobian-vector product (JVP) of fn. Given fn: Array -> Array, evaluates (fn(primal), J_fn(primal) * tangent). Propagates tangents forward through the traced equation tape in O(1) memory overhead.

Forward-mode Jacobian-vector product of fn at primal with direction tangent.

function Jvp¶

template <typename Fn >
inline JvpResult Jvp(
    Fn && fn,
    const Array & primal,
    const Array & tangent
)

Callable overload of jvp.

function Linearize¶

TENZR_APILinearizeResult Linearize(
    const std::function< Array(const Array &)> & fn,
    const Array & primal
)

Parameters:

• fn Function to linearize.
• primal Point at which to linearize.

Return: The primal output and a function mapping a tangent to its JVP.

Produces a linear approximation to fn at primal using JVP. Returns a pair (fn(primal), lin_fn), where lin_fn: tangent -> J_fn(primal) * tangent. lin_fn uses the pre-evaluated primals without re-tracing or re-evaluating non-linear functions.

Linearize fn at primal, returning its value and a tangent function.

function Linearize¶

template <typename Fn >
inline LinearizeResult Linearize(
    Fn && fn,
    const Array & primal
)

Callable overload of linearize.

function LowerNodeToTxpr¶

Txpr LowerNodeToTxpr(
    const std::shared_ptr< Node > & root
)

Parameters:

• root Root of the AST graph to lower.

Return: The equivalent Txpr program (empty when root is null).

Lower a Node AST into an equivalent Txpr SSA program.

Every VariableNode becomes a marked input (data-bound leaves are runtime arguments, never baked constants), and each op node becomes one equation in the same postorder the eager emitter uses, so executable argument order is identical to collect_parameters.

function LowerNodeToTxpr¶

Txpr LowerNodeToTxpr(
    const std::shared_ptr< Node > & root,
    std::unordered_map< uint64_t, VarId > & node_to_var
)

Parameters:

• root Root of the AST graph to lower.
• node_to_var Filled with each Node id's produced VarId.

Return: The equivalent Txpr program.

Lower a Node AST into a Txpr SSA program, also reporting the Node id -> VarId map.

function OpCodeToString¶

const char * OpCodeToString(
    OpCode op
)

Returns a stable, human-readable name for op.

Parameters:

• op Opcode to name.

Return: Stable lowercase name, or "unknown".

Short mnemonic name of an IR opcode, used for debugging and diagnostics.

function BatchTxpr¶

Txpr BatchTxpr(
    const Txpr & orig,
    int64_t batch_size,
    const std::vector< int > & in_axes ={0},
    int out_axis =0
)

Transforms an unbatched Txpr into a batched Txpr over a given batch dimension.

Parameters:

• orig Source graph.
• batch_size Batch size to introduce.
• in_axes Per-input batch axis, or -1 to leave an input unbatched.
• out_axis Batch axis for the output.

Return: The vectorized graph.

Rewrite a Txpr so every batched input gains a leading batch dimension.

function Vmap¶

template <typename Fn >
auto Vmap(
    Fn && fn,
    std::vector< int > in_axes ={0},
    int out_axis =0
)

Function transformation analogous to jax.vmap.

function Vmap¶

template <typename Fn >
auto Vmap(
    Fn && fn,
    int in_axis,
    int out_axis =0
)

Convenience overload mapping a single named input axis.

function VjpSweep¶

std::unordered_map< uint64_t, Array > VjpSweep(
    const Array & loss,
    const std::vector< Array > & targets,
    const Array & cotangent
)

Parameters:

• loss Scalar loss to differentiate.
• targets Arrays whose gradients are requested.
• cotangent Seed cotangent (typically ones).

Return: Map from target id to its gradient.

Reverse-mode sweep computing adjoints of targets from a scalar loss.