gradgen package

Submodules

gradgen.ad module

Automatic differentiation for SX expressions.

gradgen.ad.jvp(expr, wrt, tangent=None)[source]

Compute a symbolic Jacobian-vector product in forward mode.

Parameters:

  • expr (SX | SXVector) – Expression to differentiate.

  • wrt (SX | SXVector) – Scalar or vector variables with respect to which expr is differentiated.

  • tangent (SX | SXVector | float | int | list[object] | tuple[object, ...] | None) – Tangent seed matching the shape of wrt. For scalar differentiation this defaults to 1.0.

Returns:

A symbolic expression with the same shape as expr containing the forward-mode derivative in the supplied direction.

Return type:

SX | SXVector

gradgen.ad.derivative(expr, wrt, seed=1.0)[source]

Compute the forward derivative of a scalar expression.

Parameters:

  • expr (SX)

  • wrt (SX)

  • seed (SX | float | int)

Return type:

SX

gradgen.ad.vjp(expr, wrt, cotangent=None)[source]

Compute a symbolic vector-Jacobian product in reverse mode.

Parameters:

  • expr (SX | SXVector) – Output expression being differentiated.

  • wrt (SX | SXVector) – Scalar or vector variables with respect to which expr is differentiated.

  • cotangent (SX | SXVector | float | int | list[object] | tuple[object, ...] | None) – Cotangent seed matching the shape of expr. For scalar outputs this defaults to 1.0.

Returns:

A symbolic expression with the same shape as wrt containing the reverse-mode derivative in the supplied cotangent direction.

Return type:

SX | SXVector

gradgen.ad.gradient(expr, wrt, seed=1.0)[source]

Compute a reverse-mode gradient of a scalar expression.

Parameters:
Return type:

SX | SXVector

gradgen.ad.jacobian(expr, wrt)[source]

Compute a symbolic Jacobian.

The return shape follows the current no-matrix design:

  • scalar wrt scalar -> SX

  • scalar wrt vector -> SXVector

  • vector wrt scalar -> SXVector

  • vector wrt vector -> flat row-major SXVector

Parameters:
Return type:

SX | SXVector

gradgen.ad.hessian(expr, wrt)[source]

Compute a symbolic Hessian for a scalar expression.

The current return shape follows the vector-first representation:

  • scalar wrt scalar -> SX

  • scalar wrt vector -> tuple[SXVector, ...], one row per variable

Parameters:
Return type:

SX | tuple[SXVector, …]

gradgen.composed_function module

Staged function composition with packed-parameter support.

class gradgen.composed_function.ComposedGradientFunction(composed, name, simplification=None)[source]

Bases: object

Derivative kernel for a finished ComposedFunction.

The expanded derivative is the Jacobian of the composed rollout with respect to the state input.

Parameters:
composed: ComposedFunction
name: str
simplification: int | str | None
to_function(name=None)[source]

Expand this staged derivative into a regular symbolic Function.

Returns:

A symbolic function whose output is the flattened Jacobian of the finished rollout with respect to the state input.

Parameters:

name (str | None)

Return type:

Function

property nodes

Return dependency nodes for shared-helper discovery.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged derivative kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged derivative kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.composed_function.ComposedJacobianFunction(composed, name, simplification=None)[source]

Bases: object

Jacobian kernel for a finished ComposedFunction.

Parameters:
composed: ComposedFunction
name: str
simplification: int | str | None
to_function(name=None)[source]

Expand this staged Jacobian into a regular symbolic Function.

Parameters:

name (str | None)

Return type:

Function

property nodes

Return dependency nodes for shared-helper discovery.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged Jacobian kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged Jacobian kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.composed_function.ComposedJointFunction(composed, name, components, wrt_index=0, simplification=None)[source]

Bases: object

Joint kernel for a finished ComposedFunction.

Parameters:

  • composed (ComposedFunction)

  • name (str)

  • components (tuple[str, ...])

  • wrt_index (int)

  • simplification (int | str | None)

composed: ComposedFunction
name: str
components: tuple[str, ...]
wrt_index: int
simplification: int | str | None
to_function(name=None)[source]

Expand this staged joint kernel into a regular symbolic function.

Parameters:

name (str | None)

Return type:

Function

property nodes

Return dependency nodes for shared-helper discovery.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged joint kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged joint kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.composed_function.ComposedFunction(name, state_input, input_name=None, parameter_name='parameters', simplification=None, steps=(), finished=False)[source]

Bases: object

Staged composition of vector state transforms.

A ComposedFunction represents a chain such as

x -> G1(x, p1) -> G2(..., p2) -> ... -> GN(...)

while keeping the stage structure available for Rust code generation. Symbolic stage parameters are packed into one additional runtime input slice named parameter_name.

Parameters:

  • name (str)

  • state_input (SX | SXVector)

  • input_name (str | None)

  • parameter_name (str)

  • simplification (int | str | None)

  • steps (tuple[_SingleStage | _RepeatStage, ...])

  • finished (bool)

name: str
state_input: SX | SXVector
input_name: str | None
parameter_name: str
simplification: int | str | None
steps: tuple[_SingleStage | _RepeatStage, ...]
finished: bool
property parameter_size: int

Return the packed runtime parameter length.

property stage_count: int

Return the total number of concrete stage applications.

property inputs: tuple[SX | SXVector, ...]

Return the compiled symbolic inputs for this composition.

property outputs: tuple[SX | SXVector, ...]

Return the compiled symbolic outputs for this composition.

property nodes

Return dependency nodes for shared-helper discovery.

property input_names: tuple[str, ...]

Return the compiled symbolic input names.

property output_names: tuple[str, ...]

Return the compiled symbolic output names.

then(function, *, p=None)[source]

Append one explicit state-transform stage.

Parameters:

  • function (Function)

  • p (SX | SXVector | float | int | list[object] | tuple[object, ...] | None)

Return type:

ComposedFunction

chain(stages)[source]

Append a heterogeneous list of stages.

Parameters:

stages (Iterable[Function | tuple[Function, SX | SXVector | float | int | list[object] | tuple[object, ...] | None]])

Return type:

ComposedFunction

repeat(function, *, params)[source]

Append repeated applications of the same state-transform stage.

Parameters:

  • function (Function)

  • params (Iterable[SX | SXVector | float | int | list[object] | tuple[object, ...] | None])

Return type:

ComposedFunction

finish()[source]

Finalize the staged composition without adding another stage.

Returns:

A finished copy of the staged composition. The expanded function evaluates the repeated stage sequence and returns the final state.

Return type:

ComposedFunction

to_function(name=None)[source]

Expand the staged composition into a symbolic Function.

Parameters:

name (str | None)

Return type:

Function

gradient(name=None)[source]

Return a staged derivative kernel with respect to the state input.

Parameters:

name (str | None)

Return type:

ComposedGradientFunction

jacobian(name=None)[source]

Return a staged Jacobian kernel with respect to the state input.

Parameters:

name (str | None)

Return type:

ComposedJacobianFunction

joint(components, wrt_index=0, name=None, simplify_joint=None)[source]

Return a staged joint kernel for supported composed components.

Parameters:

  • components (Iterable[str])

  • wrt_index (int)

  • name (str | None)

  • simplify_joint (int | str | None)

Return type:

ComposedJointFunction

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged primal kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged primal kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

gradgen.composer module

Utilities for composing function-like stages into one pipeline.

class gradgen.composer.FunctionComposer(function)[source]

Bases: object

Build a linear composition of function-like stages.

The composer preserves staged wrappers such as ReducedFunction and BatchedFunction when Rust code is generated.

Parameters:

function (FunctionLike)

property stages: tuple[_ComposerStage, ...]

Return the configured composition stages.

feed_into(function, *, arg)[source]

Append a stage that receives the previous stage output.

Parameters:

  • function (Any)

  • arg (str)

Return type:

FunctionComposer

compose(name=None)[source]

Finalize the composer into a function-like composed pipeline.

Parameters:

name (str | None)

Return type:

FunctionComposition

class gradgen.composer.FunctionComposition(name, stages)[source]

Bases: object

A composed function-like pipeline built from staged stages.

Parameters:

  • name (str)

  • stages (tuple[_ComposerStage, ...])

name: str
stages: tuple[_ComposerStage, ...]
property input_names: tuple[str, ...]

Return the free input names of the composed pipeline.

property output_names: tuple[str, ...]

Return the output names of the composed pipeline.

property nodes

Return dependency nodes for shared-helper discovery.

to_function(name=None)[source]

Lower the composed pipeline into a symbolic Function.

Parameters:

name (str | None)

Return type:

Function

gradient(name=None)[source]

Return a gradient function for the composed pipeline.

Parameters:

name (str | None)

Return type:

Function

hessian(wrt_index=0, name=None)[source]

Return a Hessian function for the composed pipeline.

Parameters:

  • wrt_index (int)

  • name (str | None)

Return type:

Function

hvp(wrt_index=0, name=None, tangent_name=None)[source]

Return a Hessian-vector product function for the pipeline.

Parameters:

  • wrt_index (int)

  • name (str | None)

  • tangent_name (str | None)

Return type:

Function

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate Rust that calls each composed stage in sequence.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

Return type:

RustCodegenResult

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the composed pipeline.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

Return type:

RustProjectResult

gradgen.cse module

Common subexpression elimination utilities.

class gradgen.cse.CSEAssignment(name, expr, use_count)[source]

Bases: object

A named reusable intermediate expression.

Parameters:

  • name (str)

  • expr (SX)

  • use_count (int)

name: str
expr: SX
use_count: int
class gradgen.cse.CSEPlan(assignments, outputs, use_counts)[source]

Bases: object

A computation plan extracted from a symbolic DAG.

Parameters:
assignments: tuple[CSEAssignment, ...]
outputs: tuple[SX, ...]
use_counts: dict[SXNode, int]
gradgen.cse.cse(outputs, *, prefix='w', min_uses=2)[source]

Build a common-subexpression elimination plan for symbolic outputs.

Parameters:

  • outputs (Iterable[SX | SXVector]) – Scalar or vector symbolic outputs to analyze.

  • prefix (str) – Prefix used for temporary names.

  • min_uses (int) – Minimum number of uses required before an expression is promoted to a named temporary.

Returns:

A CSEPlan containing topologically ordered assignments for reusable intermediates and the flattened scalar outputs.

Return type:

CSEPlan

gradgen.function module

Symbolic function abstraction built on top of SX expressions.

class gradgen.function.Function(name, inputs, outputs, input_names=None, output_names=None)[source]

Bases: object

Symbolic multi-input, multi-output function.

A Function defines a boundary around symbolic expressions. Inputs must be symbolic leaves, while outputs can be arbitrary SX or SXVector expressions that depend on those inputs.

The class supports:

  • explicit input and output grouping

  • validation of names and symbolic inputs

  • topological ordering of the output dependency graph

  • calling with symbolic values for substitution

  • calling with numeric values for direct evaluation

Parameters:

  • name (str)

  • inputs (tuple[SX | SXVector, ...])

  • outputs (tuple[SX | SXVector, ...])

  • input_names (tuple[str, ...])

  • output_names (tuple[str, ...])

name: str
inputs: tuple[SX | SXVector, ...]
outputs: tuple[SX | SXVector, ...]
input_names: tuple[str, ...]
output_names: tuple[str, ...]
property flat_inputs: tuple[SX, ...]

Return all scalar input leaves in declaration order.

Vector inputs are flattened into their constituent scalar symbols so downstream algorithms can operate on a uniform scalar sequence.

property flat_outputs: tuple[SX, ...]

Return all scalar output leaves in declaration order.

Vector outputs are flattened into their constituent scalar symbols so downstream algorithms can operate on a uniform scalar sequence.

property nodes: tuple[SXNode, ...]

Return the expression nodes needed to compute the outputs.

Nodes are returned in topological order so that every dependency appears before the node that uses it.

cse(*, prefix='w', min_uses=2)[source]

Build a common-subexpression elimination plan for the outputs.

Parameters:

  • prefix (str) – Prefix used when naming temporary variables in the CSE plan.

  • min_uses (int) – Minimum number of times a subexpression must appear before it is extracted into a temporary.

Returns:

A CSE plan object describing the shared subexpressions in the outputs.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate Rust source code for primal function evaluation.

Parameters:

  • config – Optional backend configuration object. When omitted, a default Rust backend configuration is used.

  • function_name (str | None) – Optional Rust function name for the generated kernel. When omitted, the symbolic function name is used.

  • backend_mode (str) – Backend mode used when no explicit config is supplied.

  • scalar_type (str) – Scalar type used when no explicit config is supplied.

Returns:

A code-generation result containing the rendered Rust source and associated metadata.

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust project containing the generated primal code.

Parameters:

  • path (str) – Directory in which the generated crate should be written.

  • config – Optional backend configuration object.

  • crate_name (str | None) – Optional crate name override.

  • function_name (str | None) – Optional Rust function name for the generated kernel.

  • backend_mode (str) – Backend mode used when no explicit config is supplied.

  • scalar_type (str) – Scalar type used when no explicit config is supplied.

Returns:

A project result describing the generated crate and its supporting files.

create_rust_derivative_bundle(path, *, config=None, include_primal=True, include_jacobians=True, include_hessians=True, simplify_derivatives=None)[source]

Create a directory containing Rust projects for derivative kernels.

Parameters:

  • path (str) – Output directory for the derivative bundle.

  • config – Optional backend configuration object.

  • include_primal (bool) – When True, include the primal function crate.

  • include_jacobians (bool) – When True, include Jacobian crates.

  • include_hessians (bool) – When True, include Hessian crates.

  • simplify_derivatives (int | str | None) – Optional simplification effort applied to derivative expressions before code generation.

Returns:

A bundle result describing the generated derivative projects.

jvp(*tangent_inputs, name=None)[source]

Build a new function representing a forward-mode directional derivative.

Parameters:

  • *tangent_inputs (SX | SXVector | float | int | list[object] | tuple[object, ...]) – One tangent seed for each declared input. Seeds must match the shape of the corresponding inputs.

  • name (str | None) – Optional name for the differentiated function.

Returns:

A new Function with the same primal inputs and outputs equal to the Jacobian-vector product in the supplied tangent direction.

Return type:

Function

gradient(wrt_index=0, name=None)[source]

Build a gradient function for a scalar-output function.

Parameters:

  • wrt_index (int) – Index of the declared input block with respect to which the gradient is computed.

  • name (str | None) – Optional name for the differentiated function.

Returns:

A new Function whose single output is the gradient with respect to the selected input block.

Return type:

Function

joint(components, wrt_index=0, name=None, simplify_joint=None)[source]

Build a function that returns several requested artifacts together.

Supported component names are:

  • "f" for the primal outputs

  • "grad" for the gradient block with respect to wrt_index

  • "jf" for the Jacobian block with respect to wrt_index

  • "hessian" for the Hessian block with respect to wrt_index

  • "hvp" for the Hessian-vector product wrt wrt_index

The output order follows components exactly. If "hvp" is requested, the returned function appends one tangent input matching the selected differentiation block.

Parameters:

  • components (Iterable[str])

  • wrt_index (int)

  • name (str | None)

  • simplify_joint (int | str | None)

Return type:

Function

hvp(wrt_index=0, name=None, tangent_name=None)[source]

Build a Hessian-vector product function for a scalar-output function.

The returned function keeps the original primal inputs and appends one additional tangent input matching the selected differentiation block.

Parameters:

  • wrt_index (int)

  • name (str | None)

  • tangent_name (str | None)

Return type:

Function

vjp(*cotangent_outputs, wrt_index=None, name=None, cotangent_names=None)[source]

Build a new function representing a reverse-mode derivative.

Parameters:

  • *cotangent_outputs (SX | SXVector | float | int | list[object] | tuple[object, ...]) – One cotangent seed for each declared output. Seeds must match the shape of the corresponding outputs.

  • wrt_index (int | None) – Optional selected input block for a runtime-seeded VJP. When provided, the returned function appends cotangent inputs to the primal inputs and returns the sensitivity with respect to the selected input block only.

  • name (str | None) – Optional name for the differentiated function.

  • cotangent_names (Iterable[str] | None) – Optional names for runtime cotangent inputs.

Returns:

If explicit cotangent outputs are provided, returns a new Function with the same primal inputs and outputs equal to the vector-Jacobian product in the supplied cotangent direction, with outputs ordered like the function inputs.

If wrt_index is provided instead, returns a function with the original primal inputs plus one cotangent input per declared output, and a single output block matching the selected input.

Return type:

Function

jacobian(wrt_index=0, name=None)[source]

Build a new function representing a Jacobian block.

Parameters:

  • wrt_index (int) – Index of the declared input with respect to which the Jacobian is computed.

  • name (str | None) – Optional name for the differentiated function.

Returns:

A new Function with the same primal inputs and outputs equal to the Jacobian of each original output with respect to the selected input block.

Return type:

Function

Notes

Since full matrix types are not implemented yet, vector-output by vector-input Jacobians are returned as one flat row-major SXVector per original output.

jacobian_blocks(wrt_indices=None)[source]

Build Jacobian functions for one or more input blocks.

Parameters:

wrt_indices (Iterable[int] | None) – Optional iterable of input block indices. When omitted, Jacobians are built for every declared input block.

Returns:

A tuple of Jacobian functions, one per selected input block.

Return type:

tuple[Function, …]

vjp_blocks(wrt_indices=None)[source]

Build runtime-seeded vector-Jacobian-product functions.

Parameters:

wrt_indices (Iterable[int] | None) – Optional iterable of input block indices. When omitted, VJP functions are built for every declared input block.

Returns:

A tuple of runtime-seeded VJP functions, one per selected input block.

Return type:

tuple[Function, …]

hessian(wrt_index=0, name=None)[source]

Build a new function representing a Hessian block.

Parameters:

  • wrt_index (int) – Index of the declared input with respect to which the Hessian is computed.

  • name (str | None) – Optional name for the differentiated function.

Returns:

A new Function with the same primal inputs and outputs equal to the Hessian of the scalar output with respect to the selected input block.

Return type:

Function

Notes

Hessians are currently defined only for single scalar-output functions. For vector inputs, the Hessian is returned as a single flat SXVector in row-major order.

hessian_blocks(wrt_indices=None)[source]

Build Hessian functions for one or more input blocks.

Parameters:

wrt_indices (Iterable[int] | None) – Optional iterable of input block indices. When omitted, Hessians are built for every declared input block.

Returns:

A tuple of Hessian functions, one per selected input block.

Return type:

tuple[Function, …]

hvp_blocks(wrt_indices=None)[source]

Build Hessian-vector product functions for one or more input blocks.

Parameters:

wrt_indices (Iterable[int] | None) – Optional iterable of input block indices. When omitted, HVP functions are built for every declared input block.

Returns:

A tuple of Hessian-vector product functions, one per selected input block.

Return type:

tuple[Function, …]

simplify(max_effort='basic', name=None)[source]

Build a new function with simplified outputs.

Parameters:

  • max_effort (int | str) – Maximum simplification effort. This can be an integer pass count or one of the named presets supported by gradgen.simplify().

  • name (str | None) – Optional name for the simplified function.

Returns:

A new Function with the same inputs and simplified outputs.

Return type:

Function

gradgen.map_zip module

Loop-structured batched map/zip/reduce function abstractions.

class gradgen.map_zip.BatchedJacobianFunction(batched, wrt_index, name, simplification=None)[source]

Bases: object

Staged Jacobian wrapper for a batched function.

This object represents the Jacobian of a staged batched function without immediately expanding it back into an ordinary Function.

Parameters:

  • batched (BatchedFunction)

  • wrt_index (int)

  • name (str)

  • simplification (int | str | None)

batched

The underlying staged batched function.

Type:

gradgen.map_zip.BatchedFunction

wrt_index

The input index with respect to which the Jacobian is taken.

Type:

int

name

The symbolic name of the staged Jacobian wrapper.

Type:

str

simplification

Optional simplification effort forwarded to the expanded symbolic function.

Type:

int | str | None

batched: BatchedFunction
wrt_index: int
name: str
simplification: int | str | None
to_function(name=None)[source]

Expand this staged Jacobian into a regular symbolic function.

Parameters:

name (str | None) – Optional name for the expanded symbolic function. When not provided, the wrapper name is reused.

Returns:

A symbolic Function representing the expanded Jacobian.

Return type:

Function

property input_names: tuple[str, ...]

Return the packed input names used by the staged Jacobian.

property output_names: tuple[str, ...]

Return the generated Jacobian output names.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged Jacobian kernel.

Parameters:

  • config – Optional backend configuration object.

  • function_name (str | None) – Optional exported Rust function name.

  • backend_mode (str) – Rust backend mode to target.

  • scalar_type (str) – Rust scalar type to emit.

Returns:

The generated Rust code bundle for this staged Jacobian.

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged Jacobian kernel.

Parameters:

  • path (str) – Directory where the Rust crate should be created.

  • config – Optional backend configuration object.

  • crate_name (str | None) – Optional crate name override.

  • function_name (str | None) – Optional exported Rust function name.

  • backend_mode (str) – Rust backend mode to target.

  • scalar_type (str) – Rust scalar type to emit.

Returns:

A Rust project bundle rooted at path.

class gradgen.map_zip.BatchedFunction(function, count, name, input_sequence_names, simplification=None)[source]

Bases: object

Staged wrapper that batches a function over packed inputs.

The wrapper represents a symbolic function evaluated repeatedly over packed input sequences, producing packed outputs suitable for Rust code generation or symbolic expansion.

Parameters:

  • function (Function)

  • count (int)

  • name (str)

  • input_sequence_names (tuple[str, ...])

  • simplification (int | str | None)

function

The symbolic function to stage.

Type:

gradgen.function.Function

count

Number of repeated stages in the packed sequence.

Type:

int

name

Symbolic name of the staged wrapper.

Type:

str

input_sequence_names

Packed input sequence names.

Type:

tuple[str, …]

simplification

Optional simplification effort forwarded to the expanded symbolic function.

Type:

int | str | None

function: Function
count: int
name: str
input_sequence_names: tuple[str, ...]
simplification: int | str | None
property input_names: tuple[str, ...]

Return the packed input sequence names.

property output_names: tuple[str, ...]

Return the packed output names.

property nodes

Return dependency nodes for symbolic fallback usage.

jacobian(wrt_index=0, *, name=None)[source]

Return a staged Jacobian wrapper for one packed input.

Parameters:

  • wrt_index (int) – Input index to differentiate with respect to.

  • name (str | None) – Optional symbolic name for the staged Jacobian wrapper.

Returns:

A BatchedJacobianFunction describing the staged Jacobian.

Return type:

BatchedJacobianFunction

to_function(name=None)[source]

Expand this staged batching into a regular symbolic function.

Parameters:

name (str | None) – Optional name for the expanded symbolic function. When not provided, the wrapper name is reused.

Returns:

A symbolic Function representing the expanded staged batch.

Return type:

Function

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged batched kernel.

Parameters:

  • config – Optional backend configuration object.

  • function_name (str | None) – Optional exported Rust function name.

  • backend_mode (str) – Rust backend mode to target.

  • scalar_type (str) – Rust scalar type to emit.

Returns:

The generated Rust code bundle for this staged batch.

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged batched kernel.

Parameters:

  • path (str) – Directory where the Rust crate should be created.

  • config – Optional backend configuration object.

  • crate_name (str | None) – Optional crate name override.

  • function_name (str | None) – Optional exported Rust function name.

  • backend_mode (str) – Rust backend mode to target.

  • scalar_type (str) – Rust scalar type to emit.

Returns:

A Rust project bundle rooted at path.

class gradgen.map_zip.ReducedFunction(function, count, name, accumulator_input_name, input_name, output_name, simplification=None)[source]

Bases: object

Staged wrapper that reduces a packed input sequence.

The wrapper represents a repeated left-fold reduction over a packed input sequence and is designed for symbolic expansion or Rust generation.

Parameters:

  • function (Function)

  • count (int)

  • name (str)

  • accumulator_input_name (str)

  • input_name (str)

  • output_name (str)

  • simplification (int | str | None)

function

The binary stage function being reduced.

Type:

gradgen.function.Function

count

Number of reduction stages.

Type:

int

name

Symbolic name of the staged reducer.

Type:

str

accumulator_input_name

Name of the accumulator input sequence.

Type:

str

input_name

Name of the packed sequence input.

Type:

str

output_name

Name of the final reduced output.

Type:

str

simplification

Optional simplification effort forwarded to the expanded symbolic function.

Type:

int | str | None

function: Function
count: int
name: str
accumulator_input_name: str
input_name: str
output_name: str
simplification: int | str | None
property input_names: tuple[str, ...]

(acc0, x_seq).

Type:

Return packed reduce input names

property output_names: tuple[str, ...]

Return the reduction output names.

property nodes

Return dependency nodes for symbolic fallback usage.

to_function(name=None)[source]

Expand this staged reduce into a regular symbolic function.

Parameters:

name (str | None) – Optional name for the expanded symbolic function. When not provided, the wrapper name is reused.

Returns:

A symbolic Function representing the expanded reduction.

Return type:

Function

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged reduce kernel.

Parameters:

  • config – Optional backend configuration object.

  • function_name (str | None) – Optional exported Rust function name.

  • backend_mode (str) – Rust backend mode to target.

  • scalar_type (str) – Rust scalar type to emit.

Returns:

The generated Rust code bundle for this staged reduction.

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged reduce kernel.

Parameters:

  • path (str) – Directory where the Rust crate should be created.

  • config – Optional backend configuration object.

  • crate_name (str | None) – Optional crate name override.

  • function_name (str | None) – Optional exported Rust function name.

  • backend_mode (str) – Rust backend mode to target.

  • scalar_type (str) – Rust scalar type to emit.

Returns:

A Rust project bundle rooted at path.

gradgen.map_zip.map_function(function, count, *, input_name=None, name=None, simplification=None)[source]

Return a staged mapped function for a unary function.

Parameters:

  • function (Function) – The symbolic function to stage. It must accept exactly one input argument.

  • count (int) – Number of repeated applications in the packed sequence.

  • input_name (str | None) – Optional name for the packed input sequence. When not provided, the name defaults to "<input_name>_seq".

  • name (str | None) – Optional name for the staged mapped function. When not provided, the generated name defaults to "<function.name>_map".

  • simplification (int | str | None) – Optional simplification effort passed through to the expanded symbolic function.

Returns:

A BatchedFunction representing the staged batched function.

Return type:

BatchedFunction

gradgen.map_zip.zip_function(function, count, *, input_names=None, name=None, simplification=None)[source]

Return a staged batched function for a multi-input function.

This helper creates a loop-structured wrapper around function so it can be evaluated repeatedly over packed input sequences. The returned object is a BatchedFunction, which can later be expanded back into a regular symbolic Function or used for Rust code generation.

Parameters:

  • function (Function) – The symbolic function to stage. It may accept one or more inputs, and each input can be either a scalar SX or an SXVector.

  • count (int) – The number of times the function should be applied over the packed input sequences.

  • input_names (tuple[str, ...] | list[str] | None) – Optional names for the packed input sequences. When not provided, each input name defaults to "<input_name>_seq".

  • name (str | None) – Optional name for the staged batched function. When not provided, the generated name defaults to "<function.name>_zip".

  • simplification (int | str | None) – Optional simplification effort passed through to the generated symbolic function when the staged wrapper is expanded.

Returns:

A BatchedFunction describing the staged batched version of function.

Return type:

BatchedFunction

Example

>>> from gradgen import Function, SX
>>> x = SX.sym("x")
>>> y = SX.sym("y")
>>> f = Function("add", (x, y), (x + y,))
>>> batched = zip_function(f, 4)
>>> batched.name
'add_zip'
>>> batched.input_names
('x_seq', 'y_seq')
gradgen.map_zip.reduce_function(function, count, *, accumulator_input_name=None, input_name=None, output_name=None, name=None, simplification=None)[source]

Return a staged left-fold reduction for a binary stage function.

Parameters:

  • function (Function) – The symbolic function to stage. It must accept exactly two inputs and produce exactly one output.

  • count (int) – Number of repeated reduction stages.

  • accumulator_input_name (str | None) – Optional name for the accumulator input. When not provided, the name defaults to "<input0>0".

  • input_name (str | None) – Optional name for the packed sequence input. When not provided, the name defaults to "<input1>_seq".

  • output_name (str | None) – Optional name for the final reduced output. When not provided, the name defaults to the stage function’s output name.

  • name (str | None) – Optional name for the staged reduction. When not provided, the generated name defaults to "<function.name>_reduce".

  • simplification (int | str | None) – Optional simplification effort passed through to the expanded symbolic function.

Returns:

A ReducedFunction describing the staged reduction.

Return type:

ReducedFunction

gradgen.simplify module

Expression simplification utilities.

gradgen.simplify.simplify(value, max_effort='basic')[source]

Simplify a symbolic expression or function.

Parameters:

  • value (Any) – Expression-like value or Function to simplify.

  • max_effort (int | str) – Maximum rewrite effort. This can be an integer pass count or one of "none", "basic", "medium", "high", or "max".

Returns:

A simplified value of the same high-level shape.

Return type:

Any

gradgen.single_shooting module

Loop-structured single-shooting optimal-control abstractions.

class gradgen.single_shooting.SingleShootingBundle(include_cost=False, include_gradient=False, include_hvp=False, include_states=False)[source]

Bases: object

Requested outputs for a joint single-shooting kernel.

Parameters:

  • include_cost (bool)

  • include_gradient (bool)

  • include_hvp (bool)

  • include_states (bool)

include_cost: bool
include_gradient: bool
include_hvp: bool
include_states: bool
add_cost()[source]

Include the total cost in the joint kernel outputs.

Return type:

SingleShootingBundle

add_gradient()[source]

Include the gradient with respect to the packed control sequence.

Return type:

SingleShootingBundle

add_hvp()[source]

Include the HVP with respect to the packed control sequence.

Return type:

SingleShootingBundle

add_rollout_states()[source]

Include the packed rollout state trajectory.

Return type:

SingleShootingBundle

class gradgen.single_shooting.SingleShootingPrimalFunction(problem, name, include_states=False, simplification=None)[source]

Bases: object

Primal single-shooting cost kernel.

Parameters:
problem: SingleShootingProblem
name: str
include_states: bool
simplification: int | str | None
to_function(name=None)[source]

Expand this staged kernel into a symbolic Function.

Parameters:

name (str | None)

Return type:

Function

property nodes

Return dependency nodes for shared helper discovery.

property input_names: tuple[str, ...]

Return the exposed input names.

property output_names: tuple[str, ...]

Return the exposed output names.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged primal kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged primal kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.single_shooting.SingleShootingGradientFunction(problem, name, include_states=False, simplification=None)[source]

Bases: object

Gradient kernel for a single-shooting optimal-control problem.

Parameters:
problem: SingleShootingProblem
name: str
include_states: bool
simplification: int | str | None
to_function(name=None)[source]

Expand this staged gradient into a regular symbolic Function.

Parameters:

name (str | None)

Return type:

Function

property nodes

Return dependency nodes for shared helper discovery.

property input_names: tuple[str, ...]

Return the exposed input names.

property output_names: tuple[str, ...]

Return the exposed output names.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged gradient kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged gradient kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.single_shooting.SingleShootingHvpFunction(problem, name, include_states=False, simplification=None)[source]

Bases: object

HVP kernel for a single-shooting optimal-control problem.

Parameters:
problem: SingleShootingProblem
name: str
include_states: bool
simplification: int | str | None
to_function(name=None)[source]

Expand this staged HVP kernel into a symbolic Function.

Parameters:

name (str | None)

Return type:

Function

property nodes

Return dependency nodes for shared helper discovery.

property input_names: tuple[str, ...]

Return the exposed input names.

property output_names: tuple[str, ...]

Return the exposed output names.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged HVP kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged HVP kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.single_shooting.SingleShootingJointFunction(problem, bundle, name, simplification=None)[source]

Bases: object

Joint cost/gradient/HVP/state kernel for a single-shooting problem.

Parameters:
problem: SingleShootingProblem
bundle: SingleShootingBundle
name: str
simplification: int | str | None
to_function(name=None)[source]

Expand this staged joint kernel into a symbolic Function.

Parameters:

name (str | None)

Return type:

Function

property nodes

Return dependency nodes for shared helper discovery.

property input_names: tuple[str, ...]

Return the exposed input names.

property output_names: tuple[str, ...]

Return the exposed output names.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged joint kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged joint kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.single_shooting.SingleShootingProblem(name, horizon=None, dynamics=None, stage_cost=None, terminal_cost=None, initial_state_name='x0', control_sequence_name='U', parameter_name='p', simplification=None, stage_penalty=None, terminal_penalty=None, penalty_weight=None)[source]

Bases: object

Deterministic single-shooting optimal-control problem.

The problem represents a fixed-horizon rollout with dynamics x_next = f(x, u, p) and total cost sum_k ell(x_k, u_k, p) + V_f(x_N, p). Optional vector-valued penalty residuals may be supplied to augment this cost as c / 2 * ||q(x_k, u_k, p)||_2^2 at each stage and c / 2 * ||q_N(x_N, p)||_2^2 at the terminal state.

Parameters:

  • name (str) – Name used for expanded functions and generated kernels.

  • horizon (int | None) – Optional positive rollout horizon. It may be supplied in the constructor or later with with_horizon().

  • dynamics (Function | None) – Optional function accepting (x, u, p) and returning the next state with the same shape as x. It may be supplied in the constructor or later with with_dynamics().

  • stage_cost (Function | None) – Optional scalar function accepting (x, u, p) and returning ell(x, u, p). It may be supplied in the constructor or later with with_stage_cost().

  • terminal_cost (Function | None) – Optional scalar function accepting (x, p) and returning V_f(x, p). It may be supplied in the constructor or later with with_terminal_cost().

  • initial_state_name (str) – Runtime input name for the initial state.

  • control_sequence_name (str) – Runtime input name for the packed controls.

  • parameter_name (str) – Runtime input name for the shared parameters.

  • simplification (int | str | None) – Optional simplification effort applied to expanded functions and derivative helper kernels.

  • stage_penalty (Function | None) – Optional vector-valued residual function accepting (x, u, p) and returning q(x, u, p).

  • terminal_penalty (Function | None) – Optional vector-valued residual function accepting (x, p) and returning q_N(x, p).

  • penalty_weight (float | SX | None) – Optional scalar c multiplying both squared residual norms. Pass a numeric value to bake c into generated code, or pass an SX symbol such as SX.sym("c") to expose c as a runtime scalar input.

name: str
horizon: int | None
dynamics: Function | None
stage_cost: Function | None
terminal_cost: Function | None
initial_state_name: str
control_sequence_name: str
parameter_name: str
simplification: int | str | None
stage_penalty: Function | None
terminal_penalty: Function | None
penalty_weight: float | SX | None
with_horizon(horizon)[source]

Return a copy configured with a positive rollout horizon.

Parameters:

horizon (int) – Number of control intervals in the single-shooting rollout. It must be a positive integer.

Returns:

A new SingleShootingProblem with the requested horizon.

Return type:

SingleShootingProblem

with_dynamics(dynamics)[source]

Return a copy configured with the dynamics function.

Parameters:

dynamics (Function) – Function accepting (x, u, p) and returning the next state with the same shape as x.

Returns:

A new SingleShootingProblem using dynamics.

Return type:

SingleShootingProblem

with_stage_cost(stage_cost)[source]

Return a copy configured with the scalar stage cost.

Parameters:

stage_cost (Function) – Function accepting (x, u, p) and returning one scalar output.

Returns:

A new SingleShootingProblem using stage_cost.

Return type:

SingleShootingProblem

with_terminal_cost(terminal_cost)[source]

Return a copy configured with the scalar terminal cost.

Parameters:

terminal_cost (Function) – Function accepting (x, p) and returning one scalar output.

Returns:

A new SingleShootingProblem using terminal_cost.

Return type:

SingleShootingProblem

with_costs(stage_cost, terminal_cost)[source]

Return a copy configured with stage and terminal costs.

Parameters:

  • stage_cost (Function) – Scalar stage cost accepting (x, u, p).

  • terminal_cost (Function) – Scalar terminal cost accepting (x, p).

Returns:

A new SingleShootingProblem using both costs.

Return type:

SingleShootingProblem

with_penalties(stage_penalty, terminal_penalty, penalty_weight)[source]

Return a copy configured with residual penalties.

Parameters:

  • stage_penalty (Function) – Vector or scalar residual accepting (x, u, p).

  • terminal_penalty (Function) – Vector or scalar residual accepting (x, p).

  • penalty_weight (float | SX) – Numeric penalty weight, or an SX symbol such as SX.sym("c") to expose the weight as a runtime input.

Returns:

A new SingleShootingProblem with residual penalties.

Return type:

SingleShootingProblem

with_input_names(*, initial_state_name=None, control_sequence_name=None, parameter_name=None)[source]

Return a copy with customized generated input names.

Parameters:

  • initial_state_name (str | None) – Optional runtime input name for the initial state.

  • control_sequence_name (str | None) – Optional runtime input name for the packed control sequence.

  • parameter_name (str | None) – Optional runtime input name for shared parameters.

Returns:

A new SingleShootingProblem with updated input names.

Return type:

SingleShootingProblem

with_simplification(simplification)[source]

Return a copy with the requested simplification effort.

Parameters:

simplification (int | str | None) – Simplification effort used for expanded functions and derivative helper kernels.

Returns:

A new SingleShootingProblem with the simplification setting.

Return type:

SingleShootingProblem

property state_size: int

Return the state dimension.

property control_size: int

Return the per-stage control dimension.

property parameter_size: int

Return the shared parameter-vector dimension.

property has_runtime_penalty_weight: bool

Return whether c is exposed as a runtime scalar input.

property penalty_weight_name: str | None

Return the runtime input name for symbolic c.

property input_names: tuple[str, ...]

Return the exposed runtime input names.

property output_names: tuple[str, ...]

Return the default primal output names.

property inputs: tuple[SX | SXVector, ...]

Return compiled symbolic inputs.

property outputs: tuple[SX | SXVector, ...]

Return compiled symbolic outputs for the primal cost kernel.

property nodes

Return dependency nodes for shared helper discovery.

primal(*, include_states=False, name=None)[source]

Return a staged primal kernel source.

Parameters:

  • include_states (bool)

  • name (str | None)

Return type:

SingleShootingProblem | SingleShootingPrimalFunction

gradient(*, include_states=False, name=None)[source]

Return a staged gradient kernel source.

Parameters:

  • include_states (bool)

  • name (str | None)

Return type:

SingleShootingGradientFunction

hvp(*, include_states=False, name=None)[source]

Return a staged Hessian-vector-product kernel source.

Parameters:

  • include_states (bool)

  • name (str | None)

Return type:

SingleShootingHvpFunction

joint(bundle, *, name=None)[source]

Return a staged joint kernel source.

Parameters:
Return type:

SingleShootingJointFunction

to_function(*, include_states=False, name=None)[source]

Expand the total-cost kernel into a symbolic Function.

Parameters:

  • include_states (bool)

  • name (str | None)

Return type:

Function

stage_total_cost_function()[source]

Return the scalar stage cost including residual penalties.

The returned function has the same (x, u, p) signature as stage_cost when penalty_weight is numeric, and (x, u, p, c) when penalty_weight is symbolic. When stage_penalty is present its output is squared, summed, multiplied by penalty_weight / 2, and added to the base stage cost.

Returns:

A scalar Function for the effective stage contribution used by primal, gradient, HVP, and Rust code-generation paths.

Return type:

Function

terminal_total_cost_function()[source]

Return the scalar terminal cost including residual penalties.

The returned function has the same (x, p) signature as terminal_cost when penalty_weight is numeric, and (x, p, c) when penalty_weight is symbolic. When terminal_penalty is present its output is squared, summed, multiplied by penalty_weight / 2, and added to the base terminal cost.

Returns:

A scalar Function for the effective terminal contribution used by primal, gradient, HVP, and Rust code-generation paths.

Return type:

Function

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the primal total-cost kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the total-cost kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

gradgen.squared_distance_to_set module

Helpers for projection-backed half-squared distance primitives.

class gradgen.squared_distance_to_set.SquaredDistanceToSet(name, _sq_distance=None, _projection=None, _sq_distance_function=None, _projection_function=None, _rust_sq_distance=None, _rust_projection=None, _input_dimension=None, _input_name='x', _function_cache=None)[source]

Bases: object

Represent a projection-backed half-squared distance primitive.

This helper lets callers define a scalar-valued primitive whose gradient is derived from a projection callback. The represented scalar quantity is expected to be the half-squared Euclidean distance to a nonempty closed convex set:

\[\tfrac12 \operatorname{dist}_C(x)^2.\]

With that convention, the gradient satisfies

\[\nabla \tfrac12 \operatorname{dist}_C(x)^2 = x - \Pi_C(x).\]

Instances are configured fluently and can then be called with symbolic vectors to produce SX expressions that participate in regular gradgen composition and Rust code generation.

Parameters:

  • name (str) – Public identifier used for the registered primitive.

  • _sq_distance (Callable[[object], object] | None)

  • _projection (Callable[[object], object] | None)

  • _sq_distance_function (Function | None)

  • _projection_function (Function | None)

  • _rust_sq_distance (str | None)

  • _rust_projection (str | None)

  • _input_dimension (int | None)

  • _input_name (str)

  • _function_cache (Function | None)

Example

>>> from gradgen import SXVector, SquaredDistanceToSet
>>> distance = (
...     SquaredDistanceToSet(name="dist_to_axis")
...     .with_sq_distance(lambda x: 0.5 * x[1] * x[1])
...     .with_projection(lambda x: (x[0], 0.0))
... )
>>> x = SXVector.sym("x", 2)
>>> expr = distance(x)
>>> expr.op
'custom_vector'
name: str
with_sq_distance(callback)[source]

Attach the half-squared distance callback.

Parameters:

callback (Callable[[object], object]) – Callable that evaluates the represented scalar half-squared distance. The callback must accept either a symbolic SXVector or a numeric tuple and return a scalar.

Returns:

self to support fluent configuration.

Return type:

SquaredDistanceToSet

with_sq_distance_function(function)[source]

Attach a symbolic half-squared distance function.

Parameters:

function (Function) – A gradgen Function with exactly one vector input block and one scalar output block representing the half- squared distance.

Returns:

self to support fluent configuration.

Return type:

SquaredDistanceToSet

with_projection(callback)[source]

Attach the projection callback used to build gradients.

Parameters:

callback (Callable[[object], object]) – Callable implementing the projection map Pi_C. The callback must accept either a symbolic SXVector or a numeric tuple and return a vector-like value of matching length.

Returns:

self to support fluent configuration.

Return type:

SquaredDistanceToSet

with_projection_function(function)[source]

Attach a symbolic projection function.

Parameters:

function (Function) – A gradgen Function with exactly one vector input block and one vector output block representing the projection map.

Returns:

self to support fluent configuration.

Return type:

SquaredDistanceToSet

with_rust_sq_distance(snippet)[source]

Attach the Rust primal helper snippet.

Parameters:

snippet (str) – Rust source defining the registered primal helper with signature fn <name>(x: &[T], w: &[T]) -> T.

Returns:

self to support fluent configuration.

Return type:

SquaredDistanceToSet

with_rust_projection(snippet)[source]

Attach the Rust projection helper snippet.

Parameters:

snippet (str) – Rust source defining a projection helper with signature fn <name>_projection(x: &[T], out: &mut [T]).

Returns:

self to support fluent configuration.

Return type:

SquaredDistanceToSet

classmethod euclidean_ball(*, name, center, radius, input_name=None)[source]

Construct a half-squared distance to a Euclidean ball.

Parameters:

  • name (str) – Public identifier for the constructed distance object.

  • center (Sequence[object]) – Ball center as a sequence of scalar values.

  • radius (float | int) – Ball radius. The radius must be strictly positive.

  • input_name (str | None) – Optional symbolic input name used by the generated internal function. When omitted, "x" is used.

Returns:

A SquaredDistanceToSet configured for the Euclidean ball {x : ||x - center||_2 <= radius}.

Raises:

ValueError – If the radius is not strictly positive or the provided center is empty.

Return type:

SquaredDistanceToSet

classmethod infinity_ball(*, name, center, radius, input_name=None)[source]

Construct a half-squared distance to an infinity-norm ball.

Parameters:

  • name (str) – Public identifier for the constructed distance object.

  • center (Sequence[object]) – Ball center as a sequence of scalar values.

  • radius (float | int) – Ball radius. The radius must be strictly positive.

  • input_name (str | None) – Optional symbolic input name used by the generated internal function. When omitted, "x" is used.

Returns:

A SquaredDistanceToSet configured for the infinity ball {x : ||x - center||_∞ <= radius}.

Raises:

ValueError – If the radius is not strictly positive or the provided center is empty.

Return type:

SquaredDistanceToSet

classmethod rectangle(*, name, xmin, xmax, input_name=None)[source]

Construct a half-squared distance to a rectangle.

Parameters:

  • name (str) – Public identifier for the constructed distance object.

  • xmin (Sequence[object]) – Lower bounds for each coordinate. Individual entries may be -inf.

  • xmax (Sequence[object]) – Upper bounds for each coordinate. Individual entries may be +inf.

  • input_name (str | None) – Optional symbolic input name used by the generated internal function. When omitted, "x" is used.

Returns:

A SquaredDistanceToSet configured for the rectangle {x : xmin <= x <= xmax}.

Raises:

ValueError – If the bounds have incompatible lengths, any lower bound exceeds its upper bound, or a bound is not a supported extended real number.

Return type:

SquaredDistanceToSet

classmethod second_order_cone(*, name, alpha, dimension, input_name=None)[source]

Construct a half-squared distance to a second-order cone.

The supported cone is

\[C_\alpha = \{x = (y, t) : \lVert y \rVert_2 \leq \alpha t\},\]

where x has length dimension, y collects the first dimension - 1 entries, and t is the last entry.

This constructor is Rust-only: it provides custom Rust helpers for the half-squared distance and the cone projection, but does not expose a Python numeric evaluation callback or symbolic helper function.

Parameters:

  • name (str) – Public identifier for the constructed distance object.

  • alpha (float | int) – Positive cone scaling parameter.

  • dimension (int) – Total dimension of x = (y, t). The dimension must be at least 2.

  • input_name (str | None) – Optional symbolic input name used by the generated internal function. When omitted, "x" is used.

Returns:

A SquaredDistanceToSet configured for the second-order cone {x = (y, t) : ||y||_2 <= alpha * t}.

Return type:

SquaredDistanceToSet

Example

>>> from gradgen import SXVector, SquaredDistanceToSet
>>> distance = SquaredDistanceToSet.second_order_cone(
...     name="soc_penalty",
...     alpha=2.0,
...     dimension=3,
... )
>>> x = SXVector.sym("x", 3)
>>> expr = distance(x)
Raises:

ValueError – If alpha is not strictly positive and finite, or if dimension is less than 2.

Parameters:

  • name (str)

  • alpha (float | int)

  • dimension (int)

  • input_name (str | None)

Return type:

SquaredDistanceToSet

to_function(name=None)[source]

Return a symbolic function view of the configured distance.

Parameters:

name (str | None) – Optional symbolic function name. When omitted, the configured primitive name is used.

Returns:

A Function representing the configured half-squared distance as a scalar-valued symbolic function.

Raises:

ValueError – If the input dimension is not known yet.

Return type:

Function

gradient(wrt_index=0, name=None)[source]

Return the gradient of the materialized distance function.

Parameters:

  • wrt_index (int) – Index of the input block to differentiate with respect to.

  • name (str | None) – Optional name for the returned symbolic function.

Returns:

A Function representing the gradient with respect to the selected input block.

Return type:

Function

jacobian(wrt_index=0, name=None)[source]

Return the Jacobian of the materialized distance function.

Parameters:

  • wrt_index (int) – Index of the input block to differentiate with respect to.

  • name (str | None) – Optional name for the returned symbolic function.

Returns:

A Function representing the Jacobian block with respect to the selected input block.

Return type:

Function

gradgen.sx module

Core SX scalar and vector expression types.

This module provides the first symbolic building block for the library:

  • SXNode is the canonical internal graph node

  • SX is the small user-facing wrapper around a scalar node

  • SXVector is a lightweight vector container built from SX values

  • helper functions such as sin() and sqrt() make expression construction feel natural from Python

The implementation uses structural interning, also called hash-consing. That means two expressions with the same structure reuse the same SXNode instance. This keeps the symbolic graph compact and gives us a directed acyclic graph (DAG) instead of repeatedly rebuilding identical subtrees.

class gradgen.sx.SXNode(op, args=(), name=None, metadata=(), value=None)[source]

Bases: object

Canonical scalar expression node.

Instances of this class are immutable and are intended to be created through make(), not by calling the dataclass constructor directly. make looks up a node in a global cache and reuses an existing node when the operation and operands are structurally identical.

Parameters:

  • op (str)

  • args (tuple[SXNode, ...])

  • name (str | None)

  • metadata (tuple[tuple[str, Hashable], ...])

  • value (float | None)

op

Operation code, such as "symbol", "const", "add", or "sin".

Type:

str

args

Child nodes for non-leaf expressions.

Type:

tuple[gradgen.sx.SXNode, …]

name

Symbol name for "symbol" nodes.

Type:

str | None

metadata

Optional symbol metadata for "symbol" nodes.

Type:

tuple[tuple[str, collections.abc.Hashable], …]

value

Numeric literal for "const" nodes.

Type:

float | None

op: str
args: tuple[SXNode, ...]
name: str | None
metadata: tuple[tuple[str, Hashable], ...]
value: float | None
classmethod make(op, args=(), *, name=None, metadata=None, value=None)[source]

Create or reuse a structurally identical node.

This method is the heart of the interning strategy. It first normalizes the operand order for commutative operations, then uses the resulting structure as a cache key.

Parameters:

  • op (str) – Operation code for the node being created.

  • args (tuple[SXNode, ...]) – Child nodes of the expression.

  • name (str | None) – Optional symbol name for "symbol" nodes.

  • metadata (dict[str, Hashable] | None) – Optional metadata attached to "symbol" nodes.

  • value (float | None) – Optional floating-point literal for "const" nodes.

Returns:

The unique canonical node matching the requested structure.

Return type:

SXNode

class gradgen.sx.SX(node)[source]

Bases: object

User-facing scalar symbolic expression wrapper.

SX is a lightweight immutable wrapper around an SXNode. Most user code should work with SX values rather than raw nodes. Operator overloading is implemented here so expressions such as x + y or sin(x) produce symbolic graph nodes naturally.

Parameters:

node (SXNode)

node: SXNode
classmethod sym(name, metadata=None)[source]

Create a symbolic scalar variable.

Parameters:

  • name (str) – Symbol name used to identify the variable.

  • metadata (dict[str, Hashable] | None) – Optional metadata associated with the symbol. This is recorded as part of the symbol identity but does not yet change algebraic, AD, or code-generation semantics.

Returns:

An SX wrapper around a canonical "symbol" node.

Return type:

SX

classmethod const(value)[source]

Create a scalar constant expression.

Integer values are converted to float so the symbolic layer has one numeric literal representation for now.

Parameters:

value (float | int)

Return type:

SX

property op: str

Return the operation code of the underlying node.

These accessors expose the interned node metadata without letting callers mutate the symbolic graph.

Returns:

The underlying operation code as a string.

property args: tuple[SX, ...]

Return child expressions as SX wrappers.

The returned tuple is reconstructed from the interned node state, so callers receive fresh wrappers but the underlying node graph remains shared.

Returns:

A tuple of child SX expressions in structural order.

property name: str | None

Return the symbol name when this expression is a symbol.

These accessors expose the interned node metadata without letting callers mutate the symbolic graph.

Returns:

The symbol name, or None when the expression is not a symbol.

property value: float | None

Return the numeric literal when this expression is a constant.

These accessors expose the interned node metadata without letting callers mutate the symbolic graph.

Returns:

The constant value as float, or None when the expression is not a constant.

property metadata: dict[str, Hashable]

Return symbol metadata as a plain dictionary.

Non-symbol expressions return an empty dictionary. The returned dictionary is detached from the interned node state.

sin()[source]

Return the symbolic sine of self.

Returns:

A new SX expression.

Return type:

SX

cos()[source]

Return the symbolic cosine of self.

Returns:

A new SX expression.

Return type:

SX

tan()[source]

Return the symbolic tangent of self.

Returns:

A new SX expression.

Return type:

SX

asin()[source]

Return the symbolic arcsine of self.

Returns:

A new SX expression.

Return type:

SX

acos()[source]

Return the symbolic arccosine of self.

Returns:

A new SX expression.

Return type:

SX

atan()[source]

Return the symbolic arctangent of self.

Returns:

A new SX expression.

Return type:

SX

atan2(other)[source]

Return the symbolic two-argument arctangent with quadrant handling.

The operands are symbolically coerced before the node is created.

Parameters:

  • lhs – Scalar-like left operand.

  • rhs – Scalar-like right operand.

  • other (object)

Returns:

A new SX expression representing atan2(lhs, rhs).

Return type:

SX

sinh()[source]

Return the symbolic hyperbolic sine of self.

Returns:

A new SX expression.

Return type:

SX

cosh()[source]

Return the symbolic hyperbolic cosine of self.

Returns:

A new SX expression.

Return type:

SX

tanh()[source]

Return the symbolic hyperbolic tangent of self.

Returns:

A new SX expression.

Return type:

SX

asinh()[source]

Return the symbolic inverse hyperbolic sine of self.

Returns:

A new SX expression.

Return type:

SX

acosh()[source]

Return the symbolic inverse hyperbolic cosine of self.

Returns:

A new SX expression.

Return type:

SX

atanh()[source]

Return the symbolic inverse hyperbolic tangent of self.

Returns:

A new SX expression.

Return type:

SX

exp()[source]

Return the symbolic exponential of self.

Returns:

A new SX expression.

Return type:

SX

expm1()[source]

Return the symbolic exp(x) - 1 of self.

Returns:

A new SX expression.

Return type:

SX

log()[source]

Return the symbolic natural logarithm of self.

Returns:

A new SX expression.

Return type:

SX

log1p()[source]

Return the symbolic log(1 + x) of self.

Returns:

A new SX expression.

Return type:

SX

sqrt()[source]

Return the symbolic square root of self.

Returns:

A new SX expression.

Return type:

SX

cbrt()[source]

Return the symbolic cube root of self.

Returns:

A new SX expression.

Return type:

SX

erf()[source]

Return the symbolic error function of self.

Returns:

A new SX expression.

Return type:

SX

erfc()[source]

Return the symbolic complementary error function of self.

Returns:

A new SX expression.

Return type:

SX

floor()[source]

Return the symbolic floor of self.

Returns:

A new SX expression.

Return type:

SX

ceil()[source]

Return the symbolic ceiling of self.

Returns:

A new SX expression.

Return type:

SX

round()[source]

Return the symbolic nearest integer of self.

Returns:

A new SX expression.

Return type:

SX

trunc()[source]

Return the symbolic truncation of self.

Returns:

A new SX expression.

Return type:

SX

fract()[source]

Return the symbolic fractional part of self.

Returns:

A new SX expression.

Return type:

SX

signum()[source]

Return the symbolic sign of self.

Returns:

A new SX expression.

Return type:

SX

hypot(other)[source]

Return the symbolic Euclidean norm of two scalar-like values.

The operands are symbolically coerced before the node is created.

Parameters:

  • lhs – Scalar-like left operand.

  • rhs – Scalar-like right operand.

  • other (object)

Returns:

A new SX expression representing hypot(lhs, rhs).

Return type:

SX

abs()[source]

Return the symbolic absolute value of self.

The result is a new SX expression that shares the same interned graph style as all other nodes.

Returns:

A new SX expression representing abs(self).

Return type:

SX

maximum(other)[source]

Return the symbolic maximum of two scalar-like expressions.

The operands are symbolically coerced before the node is created.

Parameters:

  • lhs – First scalar-like value to compare.

  • rhs – Second scalar-like value to compare.

  • other (object)

Returns:

A new SX expression representing max(lhs, rhs).

Return type:

SX

minimum(other)[source]

Return the symbolic minimum of two scalar-like expressions.

The operands are symbolically coerced before the node is created.

Parameters:

  • lhs – First scalar-like value to compare.

  • rhs – Second scalar-like value to compare.

  • other (object)

Returns:

A new SX expression representing min(lhs, rhs).

Return type:

SX

class gradgen.sx.SXVector(elements)[source]

Bases: object

Vector of scalar symbolic expressions.

SXVector intentionally stays small for the first iteration of the library. It is a thin immutable container around a tuple of SX expressions and provides only vector operations that map cleanly onto the current scalar graph implementation.

For now the class supports:

  • symbolic vector construction

  • indexing, iteration, and length

  • elementwise addition, subtraction, and division

  • scalar-vector multiplication

  • elementwise unary functions

  • dot products

  • common vector norms

Elementwise vector-vector multiplication is intentionally unsupported at this stage so that * remains reserved for scalar-vector multiplication.

Parameters:

elements (tuple[SX, ...])

elements: tuple[SX, ...]
classmethod sym(name, length, metadata=None)[source]

Create a symbolic vector with indexed scalar element names.

Parameters:

  • name (str) – Base name for the vector.

  • length (int) – Number of scalar elements in the vector.

  • metadata (dict[str, Hashable] | None) – Optional metadata copied onto each scalar symbol in the vector.

Returns:

An SXVector whose elements are named "{name}_{i}".

Return type:

SXVector

sin()[source]

Apply sine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

cos()[source]

Apply cosine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

tan()[source]

Apply tangent elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

asin()[source]

Apply arcsine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

acos()[source]

Apply arccosine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

atan()[source]

Apply arctangent elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

asinh()[source]

Apply inverse hyperbolic sine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

acosh()[source]

Apply inverse hyperbolic cosine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

atanh()[source]

Apply inverse hyperbolic tangent elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

sinh()[source]

Apply hyperbolic sine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

cosh()[source]

Apply hyperbolic cosine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

tanh()[source]

Apply hyperbolic tangent elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

exp()[source]

Apply exponential elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

expm1()[source]

Apply exp(x) - 1 elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

log()[source]

Apply natural logarithm elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

log1p()[source]

Apply log(1 + x) elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

sqrt()[source]

Apply square root elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

cbrt()[source]

Apply cube root elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

erf()[source]

Apply error function elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

erfc()[source]

Apply complementary error function elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

floor()[source]

Apply floor elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

ceil()[source]

Apply ceiling elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

round()[source]

Apply nearest integer elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

trunc()[source]

Apply truncation elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

fract()[source]

Apply fractional part elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

signum()[source]

Apply sign elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

abs()[source]

Apply absolute value elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

dot(other)[source]

Return the symbolic dot product of two vectors.

Both vectors must have the same length. Empty vectors return 0.0 so the inner product remains well-defined.

Parameters:

other (object) – Vector-like operand with the same length as self.

Returns:

A scalar SX expression representing the inner product.

Raises:

  • ValueError – Raised when the two vectors do not have the same

  • length. –

Return type:

SX

cross(other)[source]

Return the symbolic three-dimensional cross product.

Both operands must be three-dimensional vectors. The resulting symbolic vector is built from the standard determinant formula, so automatic differentiation and code generation work through the existing scalar operators.

Parameters:

other (object) – Vector-like operand with length three.

Returns:

An SXVector representing self x other.

Raises:

ValueError – Raised when either operand does not have length three.

Return type:

SXVector

sum()[source]

Return the sum of all vector elements.

Empty vectors return 0.0 so the reduction stays well-defined. The additive identity is returned for empty vectors.

Returns:

A scalar SX expression.

Return type:

SX

prod()[source]

Return the product of all vector elements.

Empty vectors return 1.0 so the reduction stays well-defined. The multiplicative identity is returned for empty vectors.

Returns:

A scalar SX expression.

Return type:

SX

max()[source]

Return the maximum vector element.

Empty vectors are rejected because a maximum is undefined there.

Returns:

A scalar SX expression.

Raises:

ValueError – Raised when the vector is empty.

Return type:

SX

min()[source]

Return the minimum vector element.

Empty vectors are rejected because a minimum is undefined there.

Returns:

A scalar SX expression.

Raises:

ValueError – Raised when the vector is empty.

Return type:

SX

mean()[source]

Return the arithmetic mean of the vector elements.

Empty vectors are rejected because a mean is undefined there. The mean divides the sum by the vector length.

Returns:

A scalar SX expression.

Raises:

ValueError – Raised when the vector is empty.

Return type:

SX

norm2()[source]

Return the Euclidean norm of the vector.

Empty vectors return 0.0. The Euclidean norm is the square root of the sum of squares.

Returns:

A scalar SX expression.

Return type:

SX

norm2sq()[source]

Return the squared Euclidean norm of the vector.

Empty vectors return 0.0. The squared Euclidean norm avoids the final square root.

Returns:

A scalar SX expression.

Return type:

SX

norm1()[source]

Return the 1-norm of the vector.

Empty vectors return 0.0. The 1-norm is the sum of absolute values.

Returns:

A scalar SX expression.

Return type:

SX

norm_inf()[source]

Return the infinity norm of the vector.

Empty vectors return 0.0. The infinity norm is the maximum absolute entry.

Returns:

A scalar SX expression.

Return type:

SX

norm_p(p)[source]

Return the p-norm of the vector.

The exponent p is coerced to a scalar expression and stored in the emitted node payload. Empty vectors return 0.0.

Parameters:

p (object) – Scalar-like exponent used for the norm.

Returns:

A scalar SX expression representing the p-norm.

Return type:

SX

norm_p_to_p(p)[source]

Return the p-norm of the vector raised to the power p.

The exponent p is coerced to a scalar expression and stored in the emitted node payload. Empty vectors return 0.0.

Parameters:

p (object) – Scalar-like exponent used for the norm.

Returns:

A scalar SX expression representing ||x||_p^p.

Return type:

SX

gradgen.sx.const(value)[source]

Create a scalar constant expression.

This is a convenience alias for SX.const().

Parameters:

value (float | int)

Return type:

SX

gradgen.sx.sin(expr)[source]

Return the symbolic sine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.cos(expr)[source]

Return the symbolic cosine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.tan(expr)[source]

Return the symbolic tangent of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.asin(expr)[source]

Return the symbolic arcsine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.acos(expr)[source]

Return the symbolic arccosine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.atan(expr)[source]

Return the symbolic arctangent of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.atan2(lhs, rhs)[source]

Return the symbolic two-argument arctangent with quadrant handling.

This is a convenience wrapper around the corresponding symbolic operator.

Parameters:

  • lhs (object) – Scalar-like left operand.

  • rhs (object) – Scalar-like right operand.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.sinh(expr)[source]

Return the symbolic hyperbolic sine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.cosh(expr)[source]

Return the symbolic hyperbolic cosine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.tanh(expr)[source]

Return the symbolic hyperbolic tangent of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.asinh(expr)[source]

Return the symbolic inverse hyperbolic sine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.acosh(expr)[source]

Return the symbolic inverse hyperbolic cosine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.atanh(expr)[source]

Return the symbolic inverse hyperbolic tangent of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.exp(expr)[source]

Return the symbolic exponential of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.expm1(expr)[source]

Return the symbolic exp(x) - 1 of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.log(expr)[source]

Return the symbolic natural logarithm of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.log1p(expr)[source]

Return the symbolic log(1 + x) of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.sqrt(expr)[source]

Return the symbolic square root of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.cbrt(expr)[source]

Return the symbolic cube root of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.erf(expr)[source]

Return the symbolic error function of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.erfc(expr)[source]

Return the symbolic complementary error function of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.floor(expr)[source]

Return the symbolic floor of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.ceil(expr)[source]

Return the symbolic ceiling of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.round(expr)[source]

Return the symbolic nearest integer of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.trunc(expr)[source]

Return the symbolic truncation of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.fract(expr)[source]

Return the symbolic fractional part of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.signum(expr)[source]

Return the symbolic sign of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sx.hypot(lhs, rhs)[source]

Return the symbolic Euclidean norm of two scalar-like values.

This is a convenience wrapper around the corresponding symbolic operator.

Parameters:

  • lhs (object) – Scalar-like left operand.

  • rhs (object) – Scalar-like right operand.

Returns:

A new SX expression representing hypot(lhs, rhs).

Return type:

SX

gradgen.sx.maximum(lhs, rhs)[source]

Return the symbolic maximum of two scalar-like expressions.

This is a convenience wrapper around the corresponding symbolic operator.

Parameters:

  • lhs (object) – First scalar-like value to compare.

  • rhs (object) – Second scalar-like value to compare.

Returns:

A new SX expression representing max(lhs, rhs).

Return type:

SX

gradgen.sx.minimum(lhs, rhs)[source]

Return the symbolic minimum of two scalar-like expressions.

This is a convenience wrapper around the corresponding symbolic operator.

Parameters:

  • lhs (object) – First scalar-like value to compare.

  • rhs (object) – Second scalar-like value to compare.

Returns:

A new SX expression representing min(lhs, rhs).

Return type:

SX

gradgen.sx.if_else(if_true: object, if_false: object, condition: object) → SX[source]
gradgen.sx.if_else(if_true: SXVector, if_false: SXVector, condition: object) → SXVector

Return a symbolic branch selected by condition.

The condition is treated numerically: zero selects if_false and any nonzero value selects if_true. Comparison expressions such as x >= y therefore work naturally as conditions.

Parameters:

  • if_true (object) – Expression returned when condition is true.

  • if_false (object) – Expression returned when condition is false.

  • condition (object) – Scalar-like selector expression.

Returns:

A scalar SX expression, or an SXVector when both branches are vectors of the same length.

Return type:

SX | SXVector

gradgen.sx.vector(values)[source]

Create an SXVector from an iterable of scalar-like values.

Each element is coerced individually so the resulting vector is fully symbolic.

Parameters:

values (Iterable[object]) – Iterable yielding scalar-like values.

Returns:

A new SXVector containing the coerced entries.

Return type:

SXVector

gradgen.sx.cross(lhs, rhs)[source]

Return the symbolic three-dimensional cross product.

This helper coerces both operands to SXVector values and returns their standard three-dimensional cross product.

Parameters:

  • lhs (object) – First vector-like operand with length three.

  • rhs (object) – Second vector-like operand with length three.

Returns:

An SXVector representing lhs x rhs.

Raises:

ValueError – Raised when either operand does not have length three, or when their lengths do not match.

Return type:

SXVector

gradgen.sx.matvec(matrix, x)[source]

Return the symbolic matrix-vector product for a constant matrix.

The matrix is validated as a dense rectangular numeric sequence and encoded into one matvec_component node per output row.

Parameters:

  • matrix (Sequence[Sequence[float | int]]) – Constant numeric matrix in row-major form.

  • x (object) – Symbolic vector operand with the same number of columns as matrix.

Returns:

An SXVector whose entries represent the matrix-vector product.

Raises:

ValueError – Raised when the matrix columns do not match the vector length.

Return type:

SXVector

gradgen.sx.quadform(matrix, x, is_symmetric=True)[source]

Return a symbolic quadratic form x^   op P x for a constant matrix P.

Parameters:

  • matrix (Sequence[Sequence[float | int]]) – Constant numeric matrix in row-major form.

  • x (object) – Symbolic vector operand with the same dimension as matrix.

  • is_symmetric (bool) – When True (default), matrix is treated as symmetric: the function verifies symmetry and builds a compact equivalent form that keeps diagonal terms and doubles upper-triangular cross terms (with lower-triangular entries set to zero). When False, all entries of matrix are used.

Returns:

A scalar SX expression representing x^    op P x.

Raises:

ValueError – If matrix is not square, dimensions do not match x, or is_symmetric=True and matrix is not symmetric.

Return type:

SX

gradgen.sx.bilinear_form(x, matrix, y)[source]

Return the symbolic bilinear form x^T P y for a constant matrix.

The inputs are validated for shape compatibility before a canonical node payload is emitted.

Parameters:

  • x (object) – Left symbolic vector operand.

  • matrix (Sequence[Sequence[float | int]]) – Constant numeric matrix in row-major form.

  • y (object) – Right symbolic vector operand.

Returns:

A scalar SX expression representing the bilinear form.

Raises:

ValueError – Raised when the matrix dimensions do not match the vectors.

Return type:

SX

gradgen.sx.parse_matvec_component_args(args)[source]

Decode a matvec_component node payload.

This helper reverses the compact node representation used by the code generator.

Parameters:

args (tuple[SX, ...]) – Stored SX payload entries.

Returns:

A tuple (rows, cols, row, matrix_values, x_values).

Raises:

ValueError – Raised when the payload shape is malformed.

Return type:

tuple[int, int, int, tuple[float, …], tuple[SX, …]]

gradgen.sx.parse_quadform_args(args)[source]

Decode a quadform node payload.

This helper reverses the compact node representation used by the code generator.

Parameters:

args (tuple[SX, ...]) – Stored SX payload entries.

Returns:

A tuple (size, matrix_values, x_values).

Raises:

ValueError – Raised when the payload shape is malformed.

Return type:

tuple[int, tuple[float, …], tuple[SX, …]]

gradgen.sx.parse_bilinear_form_args(args)[source]

Decode a bilinear_form node payload.

This helper reverses the compact node representation used by the code generator.

Parameters:

args (tuple[SX, ...]) – Stored SX payload entries.

Returns:

A tuple (rows, cols, matrix_values, x_values, y_values).

Raises:

ValueError – Raised when the payload shape is malformed.

Return type:

tuple[int, int, tuple[float, …], tuple[SX, …], tuple[SX, …]]

gradgen.sx.matrix_transpose(rows, cols, values)[source]

Transpose a flattened row-major matrix.

The input values are interpreted as a matrix with rows rows and cols columns.

Parameters:

  • rows (int) – Number of matrix rows in the original layout.

  • cols (int) – Number of matrix columns in the original layout.

  • values (tuple[float, ...]) – Flattened matrix values in row-major order.

Returns:

Flattened values for the transposed matrix, also in row-major order.

Return type:

tuple[float, …]

gradgen.sx.matrix_add(values_a, values_b)[source]

Add two flattened matrices elementwise.

Both matrices must have the same flattened length.

Parameters:

  • values_a (tuple[float, ...]) – First flattened matrix.

  • values_b (tuple[float, ...]) – Second flattened matrix.

Returns:

A flattened matrix whose entries are the pairwise sums.

Raises:

ValueError – Raised when the matrices do not have the same size.

Return type:

tuple[float, …]

Module contents

gradgen package.

class gradgen.CSEAssignment(name, expr, use_count)[source]

Bases: object

A named reusable intermediate expression.

Parameters:

  • name (str)

  • expr (SX)

  • use_count (int)

name: str
expr: SX
use_count: int
class gradgen.CSEPlan(assignments, outputs, use_counts)[source]

Bases: object

A computation plan extracted from a symbolic DAG.

Parameters:
assignments: tuple[CSEAssignment, ...]
outputs: tuple[SX, ...]
use_counts: dict[SXNode, int]
class gradgen.CodeGenerationBuilder(function=None, config=None, requests=(), simplification=None, functions=())[source]

Bases: object

Fluent builder for a generated Rust crate with one or more functions.

Parameters:
function: Function | ComposedFunction | ComposedGradientFunction | ComposedJacobianFunction | ComposedJointFunction | BatchedFunction | BatchedJacobianFunction | ReducedFunction | SingleShootingProblem | SingleShootingPrimalFunction | SingleShootingGradientFunction | SingleShootingHvpFunction | SingleShootingJointFunction | None
config: RustBuilderConfigLike | None
requests: tuple[_BuilderRequest, ...]
simplification: int | str | None
functions: tuple[_BuilderFunctionSpec, ...]
with_backend_config(config)[source]

Return a copy using config for generated Rust code.

Parameters:

config (RustBuilderConfigLike)

Return type:

CodeGenerationBuilder

with_simplification(max_effort)[source]

Return a copy applying max_effort simplification to generated kernels.

Parameters:

max_effort (int | str | None)

Return type:

CodeGenerationBuilder

add_primal(*, include_states=False)[source]

Include the primal function in the generated crate.

Parameters:

include_states (bool)

Return type:

CodeGenerationBuilder

add_gradient(*, wrt=None, include_states=False)[source]

Include gradient kernels for selected input blocks.

Parameters:

  • wrt (int | str | list[int | str] | tuple[int | str, ...] | None) – Optional selector for the input blocks with respect to which gradients should be generated. Each selector may be either a zero-based input index or the declared name of an input block. When omitted, gradients are generated for every input block, matching the previous behavior.

  • include_states (bool) – When True, include rollout-state outputs for single-shooting sources that support them.

Returns:

A new builder with the gradient request appended.

Return type:

CodeGenerationBuilder

Example

>>> builder = CodeGenerationBuilder(f).add_gradient(wrt=["x", "p"])
add_jacobian()[source]

Include Jacobian kernels for all input blocks.

Return type:

CodeGenerationBuilder

add_vjp()[source]

Include runtime-seeded vector-Jacobian-product kernels for input blocks.

Return type:

CodeGenerationBuilder

add_joint(bundle)[source]

Include kernels that compute bundled artifacts together.

Parameters:

bundle (FunctionBundle | SingleShootingBundle)

Return type:

CodeGenerationBuilder

add_hessian()[source]

Include Hessian kernels for scalar-output functions.

Return type:

CodeGenerationBuilder

add_hvp(*, include_states=False)[source]

Include Hessian-vector product kernels for scalar-output functions.

Parameters:

include_states (bool)

Return type:

CodeGenerationBuilder

for_function(function, configure=None)[source]

Start configuring kernels for function or apply a legacy callback.

Parameters:
Return type:

CodeGenerationBuilder | FunctionCodegenBuilder

build(path='.')[source]

Generate a single Rust crate inside path.

The argument is treated as the parent directory that will contain the generated crate. The crate directory itself is named from with_crate_name(...) when provided, or otherwise from the first generated source function.

Parameters:

path (str | Path)

class gradgen.ComposedFunction(name, state_input, input_name=None, parameter_name='parameters', simplification=None, steps=(), finished=False)[source]

Bases: object

Staged composition of vector state transforms.

A ComposedFunction represents a chain such as

x -> G1(x, p1) -> G2(..., p2) -> ... -> GN(...)

while keeping the stage structure available for Rust code generation. Symbolic stage parameters are packed into one additional runtime input slice named parameter_name.

Parameters:

  • name (str)

  • state_input (SX | SXVector)

  • input_name (str | None)

  • parameter_name (str)

  • simplification (int | str | None)

  • steps (tuple[_SingleStage | _RepeatStage, ...])

  • finished (bool)

name: str
state_input: SX | SXVector
input_name: str | None
parameter_name: str
simplification: int | str | None
steps: tuple[_SingleStage | _RepeatStage, ...]
finished: bool
property parameter_size: int

Return the packed runtime parameter length.

property stage_count: int

Return the total number of concrete stage applications.

property inputs: tuple[SX | SXVector, ...]

Return the compiled symbolic inputs for this composition.

property outputs: tuple[SX | SXVector, ...]

Return the compiled symbolic outputs for this composition.

property nodes

Return dependency nodes for shared-helper discovery.

property input_names: tuple[str, ...]

Return the compiled symbolic input names.

property output_names: tuple[str, ...]

Return the compiled symbolic output names.

then(function, *, p=None)[source]

Append one explicit state-transform stage.

Parameters:

  • function (Function)

  • p (SX | SXVector | float | int | list[object] | tuple[object, ...] | None)

Return type:

ComposedFunction

chain(stages)[source]

Append a heterogeneous list of stages.

Parameters:

stages (Iterable[Function | tuple[Function, SX | SXVector | float | int | list[object] | tuple[object, ...] | None]])

Return type:

ComposedFunction

repeat(function, *, params)[source]

Append repeated applications of the same state-transform stage.

Parameters:

  • function (Function)

  • params (Iterable[SX | SXVector | float | int | list[object] | tuple[object, ...] | None])

Return type:

ComposedFunction

finish()[source]

Finalize the staged composition without adding another stage.

Returns:

A finished copy of the staged composition. The expanded function evaluates the repeated stage sequence and returns the final state.

Return type:

ComposedFunction

to_function(name=None)[source]

Expand the staged composition into a symbolic Function.

Parameters:

name (str | None)

Return type:

Function

gradient(name=None)[source]

Return a staged derivative kernel with respect to the state input.

Parameters:

name (str | None)

Return type:

ComposedGradientFunction

jacobian(name=None)[source]

Return a staged Jacobian kernel with respect to the state input.

Parameters:

name (str | None)

Return type:

ComposedJacobianFunction

joint(components, wrt_index=0, name=None, simplify_joint=None)[source]

Return a staged joint kernel for supported composed components.

Parameters:

  • components (Iterable[str])

  • wrt_index (int)

  • name (str | None)

  • simplify_joint (int | str | None)

Return type:

ComposedJointFunction

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged primal kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged primal kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.ComposedGradientFunction(composed, name, simplification=None)[source]

Bases: object

Derivative kernel for a finished ComposedFunction.

The expanded derivative is the Jacobian of the composed rollout with respect to the state input.

Parameters:
composed: ComposedFunction
name: str
simplification: int | str | None
to_function(name=None)[source]

Expand this staged derivative into a regular symbolic Function.

Returns:

A symbolic function whose output is the flattened Jacobian of the finished rollout with respect to the state input.

Parameters:

name (str | None)

Return type:

Function

property nodes

Return dependency nodes for shared-helper discovery.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged derivative kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged derivative kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.ComposedJacobianFunction(composed, name, simplification=None)[source]

Bases: object

Jacobian kernel for a finished ComposedFunction.

Parameters:
composed: ComposedFunction
name: str
simplification: int | str | None
to_function(name=None)[source]

Expand this staged Jacobian into a regular symbolic Function.

Parameters:

name (str | None)

Return type:

Function

property nodes

Return dependency nodes for shared-helper discovery.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged Jacobian kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged Jacobian kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.ComposedJointFunction(composed, name, components, wrt_index=0, simplification=None)[source]

Bases: object

Joint kernel for a finished ComposedFunction.

Parameters:

  • composed (ComposedFunction)

  • name (str)

  • components (tuple[str, ...])

  • wrt_index (int)

  • simplification (int | str | None)

composed: ComposedFunction
name: str
components: tuple[str, ...]
wrt_index: int
simplification: int | str | None
to_function(name=None)[source]

Expand this staged joint kernel into a regular symbolic function.

Parameters:

name (str | None)

Return type:

Function

property nodes

Return dependency nodes for shared-helper discovery.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged joint kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged joint kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.Function(name, inputs, outputs, input_names=None, output_names=None)[source]

Bases: object

Symbolic multi-input, multi-output function.

A Function defines a boundary around symbolic expressions. Inputs must be symbolic leaves, while outputs can be arbitrary SX or SXVector expressions that depend on those inputs.

The class supports:

  • explicit input and output grouping

  • validation of names and symbolic inputs

  • topological ordering of the output dependency graph

  • calling with symbolic values for substitution

  • calling with numeric values for direct evaluation

Parameters:

  • name (str)

  • inputs (tuple[SX | SXVector, ...])

  • outputs (tuple[SX | SXVector, ...])

  • input_names (tuple[str, ...])

  • output_names (tuple[str, ...])

name: str
inputs: tuple[SX | SXVector, ...]
outputs: tuple[SX | SXVector, ...]
input_names: tuple[str, ...]
output_names: tuple[str, ...]
property flat_inputs: tuple[SX, ...]

Return all scalar input leaves in declaration order.

Vector inputs are flattened into their constituent scalar symbols so downstream algorithms can operate on a uniform scalar sequence.

property flat_outputs: tuple[SX, ...]

Return all scalar output leaves in declaration order.

Vector outputs are flattened into their constituent scalar symbols so downstream algorithms can operate on a uniform scalar sequence.

property nodes: tuple[SXNode, ...]

Return the expression nodes needed to compute the outputs.

Nodes are returned in topological order so that every dependency appears before the node that uses it.

cse(*, prefix='w', min_uses=2)[source]

Build a common-subexpression elimination plan for the outputs.

Parameters:

  • prefix (str) – Prefix used when naming temporary variables in the CSE plan.

  • min_uses (int) – Minimum number of times a subexpression must appear before it is extracted into a temporary.

Returns:

A CSE plan object describing the shared subexpressions in the outputs.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate Rust source code for primal function evaluation.

Parameters:

  • config – Optional backend configuration object. When omitted, a default Rust backend configuration is used.

  • function_name (str | None) – Optional Rust function name for the generated kernel. When omitted, the symbolic function name is used.

  • backend_mode (str) – Backend mode used when no explicit config is supplied.

  • scalar_type (str) – Scalar type used when no explicit config is supplied.

Returns:

A code-generation result containing the rendered Rust source and associated metadata.

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust project containing the generated primal code.

Parameters:

  • path (str) – Directory in which the generated crate should be written.

  • config – Optional backend configuration object.

  • crate_name (str | None) – Optional crate name override.

  • function_name (str | None) – Optional Rust function name for the generated kernel.

  • backend_mode (str) – Backend mode used when no explicit config is supplied.

  • scalar_type (str) – Scalar type used when no explicit config is supplied.

Returns:

A project result describing the generated crate and its supporting files.

create_rust_derivative_bundle(path, *, config=None, include_primal=True, include_jacobians=True, include_hessians=True, simplify_derivatives=None)[source]

Create a directory containing Rust projects for derivative kernels.

Parameters:

  • path (str) – Output directory for the derivative bundle.

  • config – Optional backend configuration object.

  • include_primal (bool) – When True, include the primal function crate.

  • include_jacobians (bool) – When True, include Jacobian crates.

  • include_hessians (bool) – When True, include Hessian crates.

  • simplify_derivatives (int | str | None) – Optional simplification effort applied to derivative expressions before code generation.

Returns:

A bundle result describing the generated derivative projects.

jvp(*tangent_inputs, name=None)[source]

Build a new function representing a forward-mode directional derivative.

Parameters:

  • *tangent_inputs (SX | SXVector | float | int | list[object] | tuple[object, ...]) – One tangent seed for each declared input. Seeds must match the shape of the corresponding inputs.

  • name (str | None) – Optional name for the differentiated function.

Returns:

A new Function with the same primal inputs and outputs equal to the Jacobian-vector product in the supplied tangent direction.

Return type:

Function

gradient(wrt_index=0, name=None)[source]

Build a gradient function for a scalar-output function.

Parameters:

  • wrt_index (int) – Index of the declared input block with respect to which the gradient is computed.

  • name (str | None) – Optional name for the differentiated function.

Returns:

A new Function whose single output is the gradient with respect to the selected input block.

Return type:

Function

joint(components, wrt_index=0, name=None, simplify_joint=None)[source]

Build a function that returns several requested artifacts together.

Supported component names are:

  • "f" for the primal outputs

  • "grad" for the gradient block with respect to wrt_index

  • "jf" for the Jacobian block with respect to wrt_index

  • "hessian" for the Hessian block with respect to wrt_index

  • "hvp" for the Hessian-vector product wrt wrt_index

The output order follows components exactly. If "hvp" is requested, the returned function appends one tangent input matching the selected differentiation block.

Parameters:

  • components (Iterable[str])

  • wrt_index (int)

  • name (str | None)

  • simplify_joint (int | str | None)

Return type:

Function

hvp(wrt_index=0, name=None, tangent_name=None)[source]

Build a Hessian-vector product function for a scalar-output function.

The returned function keeps the original primal inputs and appends one additional tangent input matching the selected differentiation block.

Parameters:

  • wrt_index (int)

  • name (str | None)

  • tangent_name (str | None)

Return type:

Function

vjp(*cotangent_outputs, wrt_index=None, name=None, cotangent_names=None)[source]

Build a new function representing a reverse-mode derivative.

Parameters:

  • *cotangent_outputs (SX | SXVector | float | int | list[object] | tuple[object, ...]) – One cotangent seed for each declared output. Seeds must match the shape of the corresponding outputs.

  • wrt_index (int | None) – Optional selected input block for a runtime-seeded VJP. When provided, the returned function appends cotangent inputs to the primal inputs and returns the sensitivity with respect to the selected input block only.

  • name (str | None) – Optional name for the differentiated function.

  • cotangent_names (Iterable[str] | None) – Optional names for runtime cotangent inputs.

Returns:

If explicit cotangent outputs are provided, returns a new Function with the same primal inputs and outputs equal to the vector-Jacobian product in the supplied cotangent direction, with outputs ordered like the function inputs.

If wrt_index is provided instead, returns a function with the original primal inputs plus one cotangent input per declared output, and a single output block matching the selected input.

Return type:

Function

jacobian(wrt_index=0, name=None)[source]

Build a new function representing a Jacobian block.

Parameters:

  • wrt_index (int) – Index of the declared input with respect to which the Jacobian is computed.

  • name (str | None) – Optional name for the differentiated function.

Returns:

A new Function with the same primal inputs and outputs equal to the Jacobian of each original output with respect to the selected input block.

Return type:

Function

Notes

Since full matrix types are not implemented yet, vector-output by vector-input Jacobians are returned as one flat row-major SXVector per original output.

jacobian_blocks(wrt_indices=None)[source]

Build Jacobian functions for one or more input blocks.

Parameters:

wrt_indices (Iterable[int] | None) – Optional iterable of input block indices. When omitted, Jacobians are built for every declared input block.

Returns:

A tuple of Jacobian functions, one per selected input block.

Return type:

tuple[Function, …]

vjp_blocks(wrt_indices=None)[source]

Build runtime-seeded vector-Jacobian-product functions.

Parameters:

wrt_indices (Iterable[int] | None) – Optional iterable of input block indices. When omitted, VJP functions are built for every declared input block.

Returns:

A tuple of runtime-seeded VJP functions, one per selected input block.

Return type:

tuple[Function, …]

hessian(wrt_index=0, name=None)[source]

Build a new function representing a Hessian block.

Parameters:

  • wrt_index (int) – Index of the declared input with respect to which the Hessian is computed.

  • name (str | None) – Optional name for the differentiated function.

Returns:

A new Function with the same primal inputs and outputs equal to the Hessian of the scalar output with respect to the selected input block.

Return type:

Function

Notes

Hessians are currently defined only for single scalar-output functions. For vector inputs, the Hessian is returned as a single flat SXVector in row-major order.

hessian_blocks(wrt_indices=None)[source]

Build Hessian functions for one or more input blocks.

Parameters:

wrt_indices (Iterable[int] | None) – Optional iterable of input block indices. When omitted, Hessians are built for every declared input block.

Returns:

A tuple of Hessian functions, one per selected input block.

Return type:

tuple[Function, …]

hvp_blocks(wrt_indices=None)[source]

Build Hessian-vector product functions for one or more input blocks.

Parameters:

wrt_indices (Iterable[int] | None) – Optional iterable of input block indices. When omitted, HVP functions are built for every declared input block.

Returns:

A tuple of Hessian-vector product functions, one per selected input block.

Return type:

tuple[Function, …]

simplify(max_effort='basic', name=None)[source]

Build a new function with simplified outputs.

Parameters:

  • max_effort (int | str) – Maximum simplification effort. This can be an integer pass count or one of the named presets supported by gradgen.simplify().

  • name (str | None) – Optional name for the simplified function.

Returns:

A new Function with the same inputs and simplified outputs.

Return type:

Function

class gradgen.FunctionBundle(items=())[source]

Bases: object

Fluent description of artifacts to compute together in one joint kernel.

Parameters:

items (tuple[_BundleItem, ...])

items: tuple[_BundleItem, ...]
add_f()[source]

Include the primal outputs.

Return type:

FunctionBundle

add_gradient(*, wrt)[source]

Include gradient blocks for one or more selected inputs.

Parameters:

wrt (int | str | list[int | str] | tuple[int | str, ...]) – One or more input selectors. Each selector may be either a zero-based input index or the declared name of an input block. When multiple selectors are provided, the bundle keeps one gradient block per selected input block.

Returns:

A new FunctionBundle with the requested gradient blocks appended.

Return type:

FunctionBundle

Example

>>> bundle = FunctionBundle().add_gradient(wrt=["x", "p"])
>>> [(item.kind, item.wrt_indices) for item in bundle.items]
[('grad', ('x', 'p'))]
add_jf(*, wrt)[source]

Include Jacobian blocks for one or more input indices.

Parameters:

wrt (int | list[int] | tuple[int, ...])

Return type:

FunctionBundle

add_hessian(*, wrt)[source]

Include Hessian blocks for one or more input indices.

Parameters:

wrt (int | list[int] | tuple[int, ...])

Return type:

FunctionBundle

add_hvp(*, wrt)[source]

Include Hessian-vector-product blocks for one or more input indices.

Parameters:

wrt (int | list[int] | tuple[int, ...])

Return type:

FunctionBundle

class gradgen.FunctionComposer(function)[source]

Bases: object

Build a linear composition of function-like stages.

The composer preserves staged wrappers such as ReducedFunction and BatchedFunction when Rust code is generated.

Parameters:

function (FunctionLike)

property stages: tuple[_ComposerStage, ...]

Return the configured composition stages.

feed_into(function, *, arg)[source]

Append a stage that receives the previous stage output.

Parameters:

  • function (Any)

  • arg (str)

Return type:

FunctionComposer

compose(name=None)[source]

Finalize the composer into a function-like composed pipeline.

Parameters:

name (str | None)

Return type:

FunctionComposition

class gradgen.FunctionComposition(name, stages)[source]

Bases: object

A composed function-like pipeline built from staged stages.

Parameters:

  • name (str)

  • stages (tuple[_ComposerStage, ...])

name: str
stages: tuple[_ComposerStage, ...]
property input_names: tuple[str, ...]

Return the free input names of the composed pipeline.

property output_names: tuple[str, ...]

Return the output names of the composed pipeline.

property nodes

Return dependency nodes for shared-helper discovery.

to_function(name=None)[source]

Lower the composed pipeline into a symbolic Function.

Parameters:

name (str | None)

Return type:

Function

gradient(name=None)[source]

Return a gradient function for the composed pipeline.

Parameters:

name (str | None)

Return type:

Function

hessian(wrt_index=0, name=None)[source]

Return a Hessian function for the composed pipeline.

Parameters:

  • wrt_index (int)

  • name (str | None)

Return type:

Function

hvp(wrt_index=0, name=None, tangent_name=None)[source]

Return a Hessian-vector product function for the pipeline.

Parameters:

  • wrt_index (int)

  • name (str | None)

  • tangent_name (str | None)

Return type:

Function

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate Rust that calls each composed stage in sequence.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

Return type:

RustCodegenResult

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the composed pipeline.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

Return type:

RustProjectResult

gradgen.RustBackendMode

alias of str

class gradgen.RustBackendConfig(backend_mode='std', scalar_type='f64', crate_name=None, function_name=None, header=None, emit_metadata_helpers=True, enable_python_interface=False, build_python_interface=True, build_crate=False, build_profile='release', additional_dependencies=())[source]

Bases: object

Configuration for generated Rust source and project layout.

Parameters:

  • backend_mode (str)

  • scalar_type (str)

  • crate_name (str | None)

  • function_name (str | None)

  • header (str | None)

  • emit_metadata_helpers (bool)

  • enable_python_interface (bool)

  • build_python_interface (bool)

  • build_crate (bool)

  • build_profile (str)

  • additional_dependencies (tuple[tuple[str, str | None], ...])

backend_mode: str
scalar_type: str
crate_name: str | None
function_name: str | None
header: str | None
emit_metadata_helpers: bool
enable_python_interface: bool
build_python_interface: bool
build_crate: bool
build_profile: str
additional_dependencies: tuple[tuple[str, str | None], ...]
with_backend_mode(backend_mode)[source]

Return a copy with a different Rust backend mode.

Parameters:

backend_mode (str)

Return type:

RustBackendConfig

with_scalar_type(scalar_type)[source]

Return a copy with a different generated Rust scalar type.

Parameters:

scalar_type (str)

Return type:

RustBackendConfig

with_crate_name(crate_name)[source]

Return a copy with a different generated crate name.

Parameters:

crate_name (str | None)

Return type:

RustBackendConfig

with_function_name(function_name)[source]

Return a copy with a different generated Rust function name.

Parameters:

function_name (str | None)

Return type:

RustBackendConfig

with_header(header)[source]

Return a copy with a custom Rust module header.

Parameters:

header (str | None) – Extra Rust source to insert near the top of generated src/lib.rs. The text is written after crate attributes such as #![no_std] and #![forbid(unsafe_code)]. The header is rendered as a small Jinja2 template with variables such as scalar_type, backend_mode, and math_library.

Returns:

A copy of the config with the provided header text.

Return type:

RustBackendConfig

with_emit_metadata_helpers(emit_metadata_helpers)[source]

Return a copy with metadata helper emission enabled or disabled.

Parameters:

emit_metadata_helpers (bool)

Return type:

RustBackendConfig

with_enable_python_interface(enable_python_interface=True)[source]

Return a copy with optional PyO3-based Python bindings enabled.

Parameters:

enable_python_interface (bool)

Return type:

RustBackendConfig

with_build_python_interface(build_python_interface=True)[source]

Return a copy with Python wrapper compilation enabled or disabled.

Parameters:

build_python_interface (bool)

Return type:

RustBackendConfig

with_build_crate(build_crate=True)[source]

Return a copy with crate compilation enabled or disabled.

Parameters:

build_crate (bool) – Whether Gradgen should compile the generated low-level Rust crate after writing it to disk.

Returns:

A copy of the config with crate compilation enabled or disabled.

Return type:

RustBackendConfig

with_build_profile(build_profile)[source]

Return a copy with a different Cargo build profile.

Parameters:

build_profile (str) – Cargo profile used when build_crate is enabled. Supported values are "debug" and "release".

Returns:

A copy of the config with the requested build profile.

Return type:

RustBackendConfig

with_additional_dependencies(dependencies)[source]

Return a copy with extra Cargo dependencies added.

Parameters:

dependencies (Iterable[str | tuple[str, str | None]]) – Additional dependencies to include in the generated Cargo.toml. Each item may be either a dependency name, in which case the version defaults to "*" in Cargo syntax, or a (name, version) pair.

Returns:

A copy of the config with the normalized dependency list.

Return type:

RustBackendConfig

class gradgen.RustCodegenResult(source, python_name, function_name, workspace_size, input_names, input_sizes, output_names, output_sizes, backend_mode, scalar_type, math_library)[source]

Bases: object

Generated Rust source and metadata for a symbolic function.

Parameters:

  • source (str)

  • python_name (str)

  • function_name (str)

  • workspace_size (int)

  • input_names (tuple[str, ...])

  • input_sizes (tuple[int, ...])

  • output_names (tuple[str, ...])

  • output_sizes (tuple[int, ...])

  • backend_mode (str)

  • scalar_type (str)

  • math_library (str | None)

source: str
python_name: str
function_name: str
workspace_size: int
input_names: tuple[str, ...]
input_sizes: tuple[int, ...]
output_names: tuple[str, ...]
output_sizes: tuple[int, ...]
backend_mode: str
scalar_type: str
math_library: str | None
class gradgen.RustDerivativeBundleResult(bundle_dir, primal, jacobians, hessians)[source]

Bases: object

Information about a generated directory of derivative Rust crates.

Parameters:
bundle_dir: Path
primal: RustProjectResult | None
jacobians: tuple[RustProjectResult, ...]
hessians: tuple[RustProjectResult, ...]
class gradgen.RustMultiFunctionProjectResult(project_dir, cargo_toml, readme, metadata_json, lib_rs, codegens, python_interface=None)[source]

Bases: object

Information about a generated single-crate multi-function project.

Parameters:
project_dir: Path
cargo_toml: Path
readme: Path
metadata_json: Path
lib_rs: Path
codegens: tuple[RustCodegenResult, ...]
python_interface: RustPythonInterfaceProjectResult | None
class gradgen.RustProjectResult(project_dir, cargo_toml, readme, metadata_json, lib_rs, codegen, python_interface=None)[source]

Bases: object

Information about a generated Rust project on disk.

Parameters:
project_dir: Path
cargo_toml: Path
readme: Path
metadata_json: Path
lib_rs: Path
codegen: RustCodegenResult
python_interface: RustPythonInterfaceProjectResult | None
class gradgen.RustPythonInterfaceProjectResult(project_dir, cargo_toml, pyproject, readme, lib_rs, module_name, low_level_crate_name)[source]

Bases: object

Information about a generated Python wrapper crate.

Parameters:

  • project_dir (Path)

  • cargo_toml (Path)

  • pyproject (Path)

  • readme (Path)

  • lib_rs (Path)

  • module_name (str)

  • low_level_crate_name (str)

project_dir: Path
cargo_toml: Path
pyproject: Path
readme: Path
lib_rs: Path
module_name: str
low_level_crate_name: str
gradgen.RustScalarType

alias of str

class gradgen.SX(node)[source]

Bases: object

User-facing scalar symbolic expression wrapper.

SX is a lightweight immutable wrapper around an SXNode. Most user code should work with SX values rather than raw nodes. Operator overloading is implemented here so expressions such as x + y or sin(x) produce symbolic graph nodes naturally.

Parameters:

node (SXNode)

node: SXNode
classmethod sym(name, metadata=None)[source]

Create a symbolic scalar variable.

Parameters:

  • name (str) – Symbol name used to identify the variable.

  • metadata (dict[str, Hashable] | None) – Optional metadata associated with the symbol. This is recorded as part of the symbol identity but does not yet change algebraic, AD, or code-generation semantics.

Returns:

An SX wrapper around a canonical "symbol" node.

Return type:

SX

classmethod const(value)[source]

Create a scalar constant expression.

Integer values are converted to float so the symbolic layer has one numeric literal representation for now.

Parameters:

value (float | int)

Return type:

SX

property op: str

Return the operation code of the underlying node.

These accessors expose the interned node metadata without letting callers mutate the symbolic graph.

Returns:

The underlying operation code as a string.

property args: tuple[SX, ...]

Return child expressions as SX wrappers.

The returned tuple is reconstructed from the interned node state, so callers receive fresh wrappers but the underlying node graph remains shared.

Returns:

A tuple of child SX expressions in structural order.

property name: str | None

Return the symbol name when this expression is a symbol.

These accessors expose the interned node metadata without letting callers mutate the symbolic graph.

Returns:

The symbol name, or None when the expression is not a symbol.

property value: float | None

Return the numeric literal when this expression is a constant.

These accessors expose the interned node metadata without letting callers mutate the symbolic graph.

Returns:

The constant value as float, or None when the expression is not a constant.

property metadata: dict[str, Hashable]

Return symbol metadata as a plain dictionary.

Non-symbol expressions return an empty dictionary. The returned dictionary is detached from the interned node state.

sin()[source]

Return the symbolic sine of self.

Returns:

A new SX expression.

Return type:

SX

cos()[source]

Return the symbolic cosine of self.

Returns:

A new SX expression.

Return type:

SX

tan()[source]

Return the symbolic tangent of self.

Returns:

A new SX expression.

Return type:

SX

asin()[source]

Return the symbolic arcsine of self.

Returns:

A new SX expression.

Return type:

SX

acos()[source]

Return the symbolic arccosine of self.

Returns:

A new SX expression.

Return type:

SX

atan()[source]

Return the symbolic arctangent of self.

Returns:

A new SX expression.

Return type:

SX

atan2(other)[source]

Return the symbolic two-argument arctangent with quadrant handling.

The operands are symbolically coerced before the node is created.

Parameters:

  • lhs – Scalar-like left operand.

  • rhs – Scalar-like right operand.

  • other (object)

Returns:

A new SX expression representing atan2(lhs, rhs).

Return type:

SX

sinh()[source]

Return the symbolic hyperbolic sine of self.

Returns:

A new SX expression.

Return type:

SX

cosh()[source]

Return the symbolic hyperbolic cosine of self.

Returns:

A new SX expression.

Return type:

SX

tanh()[source]

Return the symbolic hyperbolic tangent of self.

Returns:

A new SX expression.

Return type:

SX

asinh()[source]

Return the symbolic inverse hyperbolic sine of self.

Returns:

A new SX expression.

Return type:

SX

acosh()[source]

Return the symbolic inverse hyperbolic cosine of self.

Returns:

A new SX expression.

Return type:

SX

atanh()[source]

Return the symbolic inverse hyperbolic tangent of self.

Returns:

A new SX expression.

Return type:

SX

exp()[source]

Return the symbolic exponential of self.

Returns:

A new SX expression.

Return type:

SX

expm1()[source]

Return the symbolic exp(x) - 1 of self.

Returns:

A new SX expression.

Return type:

SX

log()[source]

Return the symbolic natural logarithm of self.

Returns:

A new SX expression.

Return type:

SX

log1p()[source]

Return the symbolic log(1 + x) of self.

Returns:

A new SX expression.

Return type:

SX

sqrt()[source]

Return the symbolic square root of self.

Returns:

A new SX expression.

Return type:

SX

cbrt()[source]

Return the symbolic cube root of self.

Returns:

A new SX expression.

Return type:

SX

erf()[source]

Return the symbolic error function of self.

Returns:

A new SX expression.

Return type:

SX

erfc()[source]

Return the symbolic complementary error function of self.

Returns:

A new SX expression.

Return type:

SX

floor()[source]

Return the symbolic floor of self.

Returns:

A new SX expression.

Return type:

SX

ceil()[source]

Return the symbolic ceiling of self.

Returns:

A new SX expression.

Return type:

SX

round()[source]

Return the symbolic nearest integer of self.

Returns:

A new SX expression.

Return type:

SX

trunc()[source]

Return the symbolic truncation of self.

Returns:

A new SX expression.

Return type:

SX

fract()[source]

Return the symbolic fractional part of self.

Returns:

A new SX expression.

Return type:

SX

signum()[source]

Return the symbolic sign of self.

Returns:

A new SX expression.

Return type:

SX

hypot(other)[source]

Return the symbolic Euclidean norm of two scalar-like values.

The operands are symbolically coerced before the node is created.

Parameters:

  • lhs – Scalar-like left operand.

  • rhs – Scalar-like right operand.

  • other (object)

Returns:

A new SX expression representing hypot(lhs, rhs).

Return type:

SX

abs()[source]

Return the symbolic absolute value of self.

The result is a new SX expression that shares the same interned graph style as all other nodes.

Returns:

A new SX expression representing abs(self).

Return type:

SX

maximum(other)[source]

Return the symbolic maximum of two scalar-like expressions.

The operands are symbolically coerced before the node is created.

Parameters:

  • lhs – First scalar-like value to compare.

  • rhs – Second scalar-like value to compare.

  • other (object)

Returns:

A new SX expression representing max(lhs, rhs).

Return type:

SX

minimum(other)[source]

Return the symbolic minimum of two scalar-like expressions.

The operands are symbolically coerced before the node is created.

Parameters:

  • lhs – First scalar-like value to compare.

  • rhs – Second scalar-like value to compare.

  • other (object)

Returns:

A new SX expression representing min(lhs, rhs).

Return type:

SX

class gradgen.SXNode(op, args=(), name=None, metadata=(), value=None)[source]

Bases: object

Canonical scalar expression node.

Instances of this class are immutable and are intended to be created through make(), not by calling the dataclass constructor directly. make looks up a node in a global cache and reuses an existing node when the operation and operands are structurally identical.

Parameters:

  • op (str)

  • args (tuple[SXNode, ...])

  • name (str | None)

  • metadata (tuple[tuple[str, Hashable], ...])

  • value (float | None)

op

Operation code, such as "symbol", "const", "add", or "sin".

Type:

str

args

Child nodes for non-leaf expressions.

Type:

tuple[gradgen.sx.SXNode, …]

name

Symbol name for "symbol" nodes.

Type:

str | None

metadata

Optional symbol metadata for "symbol" nodes.

Type:

tuple[tuple[str, collections.abc.Hashable], …]

value

Numeric literal for "const" nodes.

Type:

float | None

op: str
args: tuple[SXNode, ...]
name: str | None
metadata: tuple[tuple[str, Hashable], ...]
value: float | None
classmethod make(op, args=(), *, name=None, metadata=None, value=None)[source]

Create or reuse a structurally identical node.

This method is the heart of the interning strategy. It first normalizes the operand order for commutative operations, then uses the resulting structure as a cache key.

Parameters:

  • op (str) – Operation code for the node being created.

  • args (tuple[SXNode, ...]) – Child nodes of the expression.

  • name (str | None) – Optional symbol name for "symbol" nodes.

  • metadata (dict[str, Hashable] | None) – Optional metadata attached to "symbol" nodes.

  • value (float | None) – Optional floating-point literal for "const" nodes.

Returns:

The unique canonical node matching the requested structure.

Return type:

SXNode

class gradgen.SXVector(elements)[source]

Bases: object

Vector of scalar symbolic expressions.

SXVector intentionally stays small for the first iteration of the library. It is a thin immutable container around a tuple of SX expressions and provides only vector operations that map cleanly onto the current scalar graph implementation.

For now the class supports:

  • symbolic vector construction

  • indexing, iteration, and length

  • elementwise addition, subtraction, and division

  • scalar-vector multiplication

  • elementwise unary functions

  • dot products

  • common vector norms

Elementwise vector-vector multiplication is intentionally unsupported at this stage so that * remains reserved for scalar-vector multiplication.

Parameters:

elements (tuple[SX, ...])

elements: tuple[SX, ...]
classmethod sym(name, length, metadata=None)[source]

Create a symbolic vector with indexed scalar element names.

Parameters:

  • name (str) – Base name for the vector.

  • length (int) – Number of scalar elements in the vector.

  • metadata (dict[str, Hashable] | None) – Optional metadata copied onto each scalar symbol in the vector.

Returns:

An SXVector whose elements are named "{name}_{i}".

Return type:

SXVector

sin()[source]

Apply sine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

cos()[source]

Apply cosine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

tan()[source]

Apply tangent elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

asin()[source]

Apply arcsine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

acos()[source]

Apply arccosine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

atan()[source]

Apply arctangent elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

asinh()[source]

Apply inverse hyperbolic sine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

acosh()[source]

Apply inverse hyperbolic cosine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

atanh()[source]

Apply inverse hyperbolic tangent elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

sinh()[source]

Apply hyperbolic sine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

cosh()[source]

Apply hyperbolic cosine elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

tanh()[source]

Apply hyperbolic tangent elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

exp()[source]

Apply exponential elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

expm1()[source]

Apply exp(x) - 1 elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

log()[source]

Apply natural logarithm elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

log1p()[source]

Apply log(1 + x) elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

sqrt()[source]

Apply square root elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

cbrt()[source]

Apply cube root elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

erf()[source]

Apply error function elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

erfc()[source]

Apply complementary error function elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

floor()[source]

Apply floor elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

ceil()[source]

Apply ceiling elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

round()[source]

Apply nearest integer elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

trunc()[source]

Apply truncation elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

fract()[source]

Apply fractional part elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

signum()[source]

Apply sign elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

abs()[source]

Apply absolute value elementwise to every vector entry.

Returns:

A new SXVector.

Return type:

SXVector

dot(other)[source]

Return the symbolic dot product of two vectors.

Both vectors must have the same length. Empty vectors return 0.0 so the inner product remains well-defined.

Parameters:

other (object) – Vector-like operand with the same length as self.

Returns:

A scalar SX expression representing the inner product.

Raises:

  • ValueError – Raised when the two vectors do not have the same

  • length. –

Return type:

SX

cross(other)[source]

Return the symbolic three-dimensional cross product.

Both operands must be three-dimensional vectors. The resulting symbolic vector is built from the standard determinant formula, so automatic differentiation and code generation work through the existing scalar operators.

Parameters:

other (object) – Vector-like operand with length three.

Returns:

An SXVector representing self x other.

Raises:

ValueError – Raised when either operand does not have length three.

Return type:

SXVector

sum()[source]

Return the sum of all vector elements.

Empty vectors return 0.0 so the reduction stays well-defined. The additive identity is returned for empty vectors.

Returns:

A scalar SX expression.

Return type:

SX

prod()[source]

Return the product of all vector elements.

Empty vectors return 1.0 so the reduction stays well-defined. The multiplicative identity is returned for empty vectors.

Returns:

A scalar SX expression.

Return type:

SX

max()[source]

Return the maximum vector element.

Empty vectors are rejected because a maximum is undefined there.

Returns:

A scalar SX expression.

Raises:

ValueError – Raised when the vector is empty.

Return type:

SX

min()[source]

Return the minimum vector element.

Empty vectors are rejected because a minimum is undefined there.

Returns:

A scalar SX expression.

Raises:

ValueError – Raised when the vector is empty.

Return type:

SX

mean()[source]

Return the arithmetic mean of the vector elements.

Empty vectors are rejected because a mean is undefined there. The mean divides the sum by the vector length.

Returns:

A scalar SX expression.

Raises:

ValueError – Raised when the vector is empty.

Return type:

SX

norm2()[source]

Return the Euclidean norm of the vector.

Empty vectors return 0.0. The Euclidean norm is the square root of the sum of squares.

Returns:

A scalar SX expression.

Return type:

SX

norm2sq()[source]

Return the squared Euclidean norm of the vector.

Empty vectors return 0.0. The squared Euclidean norm avoids the final square root.

Returns:

A scalar SX expression.

Return type:

SX

norm1()[source]

Return the 1-norm of the vector.

Empty vectors return 0.0. The 1-norm is the sum of absolute values.

Returns:

A scalar SX expression.

Return type:

SX

norm_inf()[source]

Return the infinity norm of the vector.

Empty vectors return 0.0. The infinity norm is the maximum absolute entry.

Returns:

A scalar SX expression.

Return type:

SX

norm_p(p)[source]

Return the p-norm of the vector.

The exponent p is coerced to a scalar expression and stored in the emitted node payload. Empty vectors return 0.0.

Parameters:

p (object) – Scalar-like exponent used for the norm.

Returns:

A scalar SX expression representing the p-norm.

Return type:

SX

norm_p_to_p(p)[source]

Return the p-norm of the vector raised to the power p.

The exponent p is coerced to a scalar expression and stored in the emitted node payload. Empty vectors return 0.0.

Parameters:

p (object) – Scalar-like exponent used for the norm.

Returns:

A scalar SX expression representing ||x||_p^p.

Return type:

SX

class gradgen.SingleShootingBundle(include_cost=False, include_gradient=False, include_hvp=False, include_states=False)[source]

Bases: object

Requested outputs for a joint single-shooting kernel.

Parameters:

  • include_cost (bool)

  • include_gradient (bool)

  • include_hvp (bool)

  • include_states (bool)

include_cost: bool
include_gradient: bool
include_hvp: bool
include_states: bool
add_cost()[source]

Include the total cost in the joint kernel outputs.

Return type:

SingleShootingBundle

add_gradient()[source]

Include the gradient with respect to the packed control sequence.

Return type:

SingleShootingBundle

add_hvp()[source]

Include the HVP with respect to the packed control sequence.

Return type:

SingleShootingBundle

add_rollout_states()[source]

Include the packed rollout state trajectory.

Return type:

SingleShootingBundle

class gradgen.SingleShootingGradientFunction(problem, name, include_states=False, simplification=None)[source]

Bases: object

Gradient kernel for a single-shooting optimal-control problem.

Parameters:
problem: SingleShootingProblem
name: str
include_states: bool
simplification: int | str | None
to_function(name=None)[source]

Expand this staged gradient into a regular symbolic Function.

Parameters:

name (str | None)

Return type:

Function

property nodes

Return dependency nodes for shared helper discovery.

property input_names: tuple[str, ...]

Return the exposed input names.

property output_names: tuple[str, ...]

Return the exposed output names.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged gradient kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged gradient kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.SingleShootingHvpFunction(problem, name, include_states=False, simplification=None)[source]

Bases: object

HVP kernel for a single-shooting optimal-control problem.

Parameters:
problem: SingleShootingProblem
name: str
include_states: bool
simplification: int | str | None
to_function(name=None)[source]

Expand this staged HVP kernel into a symbolic Function.

Parameters:

name (str | None)

Return type:

Function

property nodes

Return dependency nodes for shared helper discovery.

property input_names: tuple[str, ...]

Return the exposed input names.

property output_names: tuple[str, ...]

Return the exposed output names.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged HVP kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged HVP kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.SingleShootingJointFunction(problem, bundle, name, simplification=None)[source]

Bases: object

Joint cost/gradient/HVP/state kernel for a single-shooting problem.

Parameters:
problem: SingleShootingProblem
bundle: SingleShootingBundle
name: str
simplification: int | str | None
to_function(name=None)[source]

Expand this staged joint kernel into a symbolic Function.

Parameters:

name (str | None)

Return type:

Function

property nodes

Return dependency nodes for shared helper discovery.

property input_names: tuple[str, ...]

Return the exposed input names.

property output_names: tuple[str, ...]

Return the exposed output names.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged joint kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged joint kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.SingleShootingPrimalFunction(problem, name, include_states=False, simplification=None)[source]

Bases: object

Primal single-shooting cost kernel.

Parameters:
problem: SingleShootingProblem
name: str
include_states: bool
simplification: int | str | None
to_function(name=None)[source]

Expand this staged kernel into a symbolic Function.

Parameters:

name (str | None)

Return type:

Function

property nodes

Return dependency nodes for shared helper discovery.

property input_names: tuple[str, ...]

Return the exposed input names.

property output_names: tuple[str, ...]

Return the exposed output names.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged primal kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged primal kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.SingleShootingProblem(name, horizon=None, dynamics=None, stage_cost=None, terminal_cost=None, initial_state_name='x0', control_sequence_name='U', parameter_name='p', simplification=None, stage_penalty=None, terminal_penalty=None, penalty_weight=None)[source]

Bases: object

Deterministic single-shooting optimal-control problem.

The problem represents a fixed-horizon rollout with dynamics x_next = f(x, u, p) and total cost sum_k ell(x_k, u_k, p) + V_f(x_N, p). Optional vector-valued penalty residuals may be supplied to augment this cost as c / 2 * ||q(x_k, u_k, p)||_2^2 at each stage and c / 2 * ||q_N(x_N, p)||_2^2 at the terminal state.

Parameters:

  • name (str) – Name used for expanded functions and generated kernels.

  • horizon (int | None) – Optional positive rollout horizon. It may be supplied in the constructor or later with with_horizon().

  • dynamics (Function | None) – Optional function accepting (x, u, p) and returning the next state with the same shape as x. It may be supplied in the constructor or later with with_dynamics().

  • stage_cost (Function | None) – Optional scalar function accepting (x, u, p) and returning ell(x, u, p). It may be supplied in the constructor or later with with_stage_cost().

  • terminal_cost (Function | None) – Optional scalar function accepting (x, p) and returning V_f(x, p). It may be supplied in the constructor or later with with_terminal_cost().

  • initial_state_name (str) – Runtime input name for the initial state.

  • control_sequence_name (str) – Runtime input name for the packed controls.

  • parameter_name (str) – Runtime input name for the shared parameters.

  • simplification (int | str | None) – Optional simplification effort applied to expanded functions and derivative helper kernels.

  • stage_penalty (Function | None) – Optional vector-valued residual function accepting (x, u, p) and returning q(x, u, p).

  • terminal_penalty (Function | None) – Optional vector-valued residual function accepting (x, p) and returning q_N(x, p).

  • penalty_weight (float | SX | None) – Optional scalar c multiplying both squared residual norms. Pass a numeric value to bake c into generated code, or pass an SX symbol such as SX.sym("c") to expose c as a runtime scalar input.

name: str
horizon: int | None
dynamics: Function | None
stage_cost: Function | None
terminal_cost: Function | None
initial_state_name: str
control_sequence_name: str
parameter_name: str
simplification: int | str | None
stage_penalty: Function | None
terminal_penalty: Function | None
penalty_weight: float | SX | None
with_horizon(horizon)[source]

Return a copy configured with a positive rollout horizon.

Parameters:

horizon (int) – Number of control intervals in the single-shooting rollout. It must be a positive integer.

Returns:

A new SingleShootingProblem with the requested horizon.

Return type:

SingleShootingProblem

with_dynamics(dynamics)[source]

Return a copy configured with the dynamics function.

Parameters:

dynamics (Function) – Function accepting (x, u, p) and returning the next state with the same shape as x.

Returns:

A new SingleShootingProblem using dynamics.

Return type:

SingleShootingProblem

with_stage_cost(stage_cost)[source]

Return a copy configured with the scalar stage cost.

Parameters:

stage_cost (Function) – Function accepting (x, u, p) and returning one scalar output.

Returns:

A new SingleShootingProblem using stage_cost.

Return type:

SingleShootingProblem

with_terminal_cost(terminal_cost)[source]

Return a copy configured with the scalar terminal cost.

Parameters:

terminal_cost (Function) – Function accepting (x, p) and returning one scalar output.

Returns:

A new SingleShootingProblem using terminal_cost.

Return type:

SingleShootingProblem

with_costs(stage_cost, terminal_cost)[source]

Return a copy configured with stage and terminal costs.

Parameters:

  • stage_cost (Function) – Scalar stage cost accepting (x, u, p).

  • terminal_cost (Function) – Scalar terminal cost accepting (x, p).

Returns:

A new SingleShootingProblem using both costs.

Return type:

SingleShootingProblem

with_penalties(stage_penalty, terminal_penalty, penalty_weight)[source]

Return a copy configured with residual penalties.

Parameters:

  • stage_penalty (Function) – Vector or scalar residual accepting (x, u, p).

  • terminal_penalty (Function) – Vector or scalar residual accepting (x, p).

  • penalty_weight (float | SX) – Numeric penalty weight, or an SX symbol such as SX.sym("c") to expose the weight as a runtime input.

Returns:

A new SingleShootingProblem with residual penalties.

Return type:

SingleShootingProblem

with_input_names(*, initial_state_name=None, control_sequence_name=None, parameter_name=None)[source]

Return a copy with customized generated input names.

Parameters:

  • initial_state_name (str | None) – Optional runtime input name for the initial state.

  • control_sequence_name (str | None) – Optional runtime input name for the packed control sequence.

  • parameter_name (str | None) – Optional runtime input name for shared parameters.

Returns:

A new SingleShootingProblem with updated input names.

Return type:

SingleShootingProblem

with_simplification(simplification)[source]

Return a copy with the requested simplification effort.

Parameters:

simplification (int | str | None) – Simplification effort used for expanded functions and derivative helper kernels.

Returns:

A new SingleShootingProblem with the simplification setting.

Return type:

SingleShootingProblem

property state_size: int

Return the state dimension.

property control_size: int

Return the per-stage control dimension.

property parameter_size: int

Return the shared parameter-vector dimension.

property has_runtime_penalty_weight: bool

Return whether c is exposed as a runtime scalar input.

property penalty_weight_name: str | None

Return the runtime input name for symbolic c.

property input_names: tuple[str, ...]

Return the exposed runtime input names.

property output_names: tuple[str, ...]

Return the default primal output names.

property inputs: tuple[SX | SXVector, ...]

Return compiled symbolic inputs.

property outputs: tuple[SX | SXVector, ...]

Return compiled symbolic outputs for the primal cost kernel.

property nodes

Return dependency nodes for shared helper discovery.

primal(*, include_states=False, name=None)[source]

Return a staged primal kernel source.

Parameters:

  • include_states (bool)

  • name (str | None)

Return type:

SingleShootingProblem | SingleShootingPrimalFunction

gradient(*, include_states=False, name=None)[source]

Return a staged gradient kernel source.

Parameters:

  • include_states (bool)

  • name (str | None)

Return type:

SingleShootingGradientFunction

hvp(*, include_states=False, name=None)[source]

Return a staged Hessian-vector-product kernel source.

Parameters:

  • include_states (bool)

  • name (str | None)

Return type:

SingleShootingHvpFunction

joint(bundle, *, name=None)[source]

Return a staged joint kernel source.

Parameters:
Return type:

SingleShootingJointFunction

to_function(*, include_states=False, name=None)[source]

Expand the total-cost kernel into a symbolic Function.

Parameters:

  • include_states (bool)

  • name (str | None)

Return type:

Function

stage_total_cost_function()[source]

Return the scalar stage cost including residual penalties.

The returned function has the same (x, u, p) signature as stage_cost when penalty_weight is numeric, and (x, u, p, c) when penalty_weight is symbolic. When stage_penalty is present its output is squared, summed, multiplied by penalty_weight / 2, and added to the base stage cost.

Returns:

A scalar Function for the effective stage contribution used by primal, gradient, HVP, and Rust code-generation paths.

Return type:

Function

terminal_total_cost_function()[source]

Return the scalar terminal cost including residual penalties.

The returned function has the same (x, p) signature as terminal_cost when penalty_weight is numeric, and (x, p, c) when penalty_weight is symbolic. When terminal_penalty is present its output is squared, summed, multiplied by penalty_weight / 2, and added to the base terminal cost.

Returns:

A scalar Function for the effective terminal contribution used by primal, gradient, HVP, and Rust code-generation paths.

Return type:

Function

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the primal total-cost kernel.

Parameters:

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the total-cost kernel.

Parameters:

  • path (str)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

class gradgen.SquaredDistanceToSet(name, _sq_distance=None, _projection=None, _sq_distance_function=None, _projection_function=None, _rust_sq_distance=None, _rust_projection=None, _input_dimension=None, _input_name='x', _function_cache=None)[source]

Bases: object

Represent a projection-backed half-squared distance primitive.

This helper lets callers define a scalar-valued primitive whose gradient is derived from a projection callback. The represented scalar quantity is expected to be the half-squared Euclidean distance to a nonempty closed convex set:

\[\tfrac12 \operatorname{dist}_C(x)^2.\]

With that convention, the gradient satisfies

\[\nabla \tfrac12 \operatorname{dist}_C(x)^2 = x - \Pi_C(x).\]

Instances are configured fluently and can then be called with symbolic vectors to produce SX expressions that participate in regular gradgen composition and Rust code generation.

Parameters:

  • name (str) – Public identifier used for the registered primitive.

  • _sq_distance (Callable[[object], object] | None)

  • _projection (Callable[[object], object] | None)

  • _sq_distance_function (Function | None)

  • _projection_function (Function | None)

  • _rust_sq_distance (str | None)

  • _rust_projection (str | None)

  • _input_dimension (int | None)

  • _input_name (str)

  • _function_cache (Function | None)

Example

>>> from gradgen import SXVector, SquaredDistanceToSet
>>> distance = (
...     SquaredDistanceToSet(name="dist_to_axis")
...     .with_sq_distance(lambda x: 0.5 * x[1] * x[1])
...     .with_projection(lambda x: (x[0], 0.0))
... )
>>> x = SXVector.sym("x", 2)
>>> expr = distance(x)
>>> expr.op
'custom_vector'
name: str
with_sq_distance(callback)[source]

Attach the half-squared distance callback.

Parameters:

callback (Callable[[object], object]) – Callable that evaluates the represented scalar half-squared distance. The callback must accept either a symbolic SXVector or a numeric tuple and return a scalar.

Returns:

self to support fluent configuration.

Return type:

SquaredDistanceToSet

with_sq_distance_function(function)[source]

Attach a symbolic half-squared distance function.

Parameters:

function (Function) – A gradgen Function with exactly one vector input block and one scalar output block representing the half- squared distance.

Returns:

self to support fluent configuration.

Return type:

SquaredDistanceToSet

with_projection(callback)[source]

Attach the projection callback used to build gradients.

Parameters:

callback (Callable[[object], object]) – Callable implementing the projection map Pi_C. The callback must accept either a symbolic SXVector or a numeric tuple and return a vector-like value of matching length.

Returns:

self to support fluent configuration.

Return type:

SquaredDistanceToSet

with_projection_function(function)[source]

Attach a symbolic projection function.

Parameters:

function (Function) – A gradgen Function with exactly one vector input block and one vector output block representing the projection map.

Returns:

self to support fluent configuration.

Return type:

SquaredDistanceToSet

with_rust_sq_distance(snippet)[source]

Attach the Rust primal helper snippet.

Parameters:

snippet (str) – Rust source defining the registered primal helper with signature fn <name>(x: &[T], w: &[T]) -> T.

Returns:

self to support fluent configuration.

Return type:

SquaredDistanceToSet

with_rust_projection(snippet)[source]

Attach the Rust projection helper snippet.

Parameters:

snippet (str) – Rust source defining a projection helper with signature fn <name>_projection(x: &[T], out: &mut [T]).

Returns:

self to support fluent configuration.

Return type:

SquaredDistanceToSet

classmethod euclidean_ball(*, name, center, radius, input_name=None)[source]

Construct a half-squared distance to a Euclidean ball.

Parameters:

  • name (str) – Public identifier for the constructed distance object.

  • center (Sequence[object]) – Ball center as a sequence of scalar values.

  • radius (float | int) – Ball radius. The radius must be strictly positive.

  • input_name (str | None) – Optional symbolic input name used by the generated internal function. When omitted, "x" is used.

Returns:

A SquaredDistanceToSet configured for the Euclidean ball {x : ||x - center||_2 <= radius}.

Raises:

ValueError – If the radius is not strictly positive or the provided center is empty.

Return type:

SquaredDistanceToSet

classmethod infinity_ball(*, name, center, radius, input_name=None)[source]

Construct a half-squared distance to an infinity-norm ball.

Parameters:

  • name (str) – Public identifier for the constructed distance object.

  • center (Sequence[object]) – Ball center as a sequence of scalar values.

  • radius (float | int) – Ball radius. The radius must be strictly positive.

  • input_name (str | None) – Optional symbolic input name used by the generated internal function. When omitted, "x" is used.

Returns:

A SquaredDistanceToSet configured for the infinity ball {x : ||x - center||_∞ <= radius}.

Raises:

ValueError – If the radius is not strictly positive or the provided center is empty.

Return type:

SquaredDistanceToSet

classmethod rectangle(*, name, xmin, xmax, input_name=None)[source]

Construct a half-squared distance to a rectangle.

Parameters:

  • name (str) – Public identifier for the constructed distance object.

  • xmin (Sequence[object]) – Lower bounds for each coordinate. Individual entries may be -inf.

  • xmax (Sequence[object]) – Upper bounds for each coordinate. Individual entries may be +inf.

  • input_name (str | None) – Optional symbolic input name used by the generated internal function. When omitted, "x" is used.

Returns:

A SquaredDistanceToSet configured for the rectangle {x : xmin <= x <= xmax}.

Raises:

ValueError – If the bounds have incompatible lengths, any lower bound exceeds its upper bound, or a bound is not a supported extended real number.

Return type:

SquaredDistanceToSet

classmethod second_order_cone(*, name, alpha, dimension, input_name=None)[source]

Construct a half-squared distance to a second-order cone.

The supported cone is

\[C_\alpha = \{x = (y, t) : \lVert y \rVert_2 \leq \alpha t\},\]

where x has length dimension, y collects the first dimension - 1 entries, and t is the last entry.

This constructor is Rust-only: it provides custom Rust helpers for the half-squared distance and the cone projection, but does not expose a Python numeric evaluation callback or symbolic helper function.

Parameters:

  • name (str) – Public identifier for the constructed distance object.

  • alpha (float | int) – Positive cone scaling parameter.

  • dimension (int) – Total dimension of x = (y, t). The dimension must be at least 2.

  • input_name (str | None) – Optional symbolic input name used by the generated internal function. When omitted, "x" is used.

Returns:

A SquaredDistanceToSet configured for the second-order cone {x = (y, t) : ||y||_2 <= alpha * t}.

Return type:

SquaredDistanceToSet

Example

>>> from gradgen import SXVector, SquaredDistanceToSet
>>> distance = SquaredDistanceToSet.second_order_cone(
...     name="soc_penalty",
...     alpha=2.0,
...     dimension=3,
... )
>>> x = SXVector.sym("x", 3)
>>> expr = distance(x)
Raises:

ValueError – If alpha is not strictly positive and finite, or if dimension is less than 2.

Parameters:

  • name (str)

  • alpha (float | int)

  • dimension (int)

  • input_name (str | None)

Return type:

SquaredDistanceToSet

to_function(name=None)[source]

Return a symbolic function view of the configured distance.

Parameters:

name (str | None) – Optional symbolic function name. When omitted, the configured primitive name is used.

Returns:

A Function representing the configured half-squared distance as a scalar-valued symbolic function.

Raises:

ValueError – If the input dimension is not known yet.

Return type:

Function

gradient(wrt_index=0, name=None)[source]

Return the gradient of the materialized distance function.

Parameters:

  • wrt_index (int) – Index of the input block to differentiate with respect to.

  • name (str | None) – Optional name for the returned symbolic function.

Returns:

A Function representing the gradient with respect to the selected input block.

Return type:

Function

jacobian(wrt_index=0, name=None)[source]

Return the Jacobian of the materialized distance function.

Parameters:

  • wrt_index (int) – Index of the input block to differentiate with respect to.

  • name (str | None) – Optional name for the returned symbolic function.

Returns:

A Function representing the Jacobian block with respect to the selected input block.

Return type:

Function

gradgen.acosh(expr)[source]

Return the symbolic inverse hyperbolic cosine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.acos(expr)[source]

Return the symbolic arccosine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.asinh(expr)[source]

Return the symbolic inverse hyperbolic sine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.asin(expr)[source]

Return the symbolic arcsine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.atan2(lhs, rhs)[source]

Return the symbolic two-argument arctangent with quadrant handling.

This is a convenience wrapper around the corresponding symbolic operator.

Parameters:

  • lhs (object) – Scalar-like left operand.

  • rhs (object) – Scalar-like right operand.

Returns:

A new SX expression.

Return type:

SX

gradgen.atan(expr)[source]

Return the symbolic arctangent of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.atanh(expr)[source]

Return the symbolic inverse hyperbolic tangent of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.bilinear_form(x, matrix, y)[source]

Return the symbolic bilinear form x^T P y for a constant matrix.

The inputs are validated for shape compatibility before a canonical node payload is emitted.

Parameters:

  • x (object) – Left symbolic vector operand.

  • matrix (Sequence[Sequence[float | int]]) – Constant numeric matrix in row-major form.

  • y (object) – Right symbolic vector operand.

Returns:

A scalar SX expression representing the bilinear form.

Raises:

ValueError – Raised when the matrix dimensions do not match the vectors.

Return type:

SX

gradgen.cbrt(expr)[source]

Return the symbolic cube root of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.ceil(expr)[source]

Return the symbolic ceiling of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.const(value)[source]

Create a scalar constant expression.

This is a convenience alias for SX.const().

Parameters:

value (float | int)

Return type:

SX

gradgen.cos(expr)[source]

Return the symbolic cosine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.cosh(expr)[source]

Return the symbolic hyperbolic cosine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.cross(lhs, rhs)[source]

Return the symbolic three-dimensional cross product.

This helper coerces both operands to SXVector values and returns their standard three-dimensional cross product.

Parameters:

  • lhs (object) – First vector-like operand with length three.

  • rhs (object) – Second vector-like operand with length three.

Returns:

An SXVector representing lhs x rhs.

Raises:

ValueError – Raised when either operand does not have length three, or when their lengths do not match.

Return type:

SXVector

gradgen.cse(outputs, *, prefix='w', min_uses=2)[source]

Build a common-subexpression elimination plan for symbolic outputs.

Parameters:

  • outputs (Iterable[SX | SXVector]) – Scalar or vector symbolic outputs to analyze.

  • prefix (str) – Prefix used for temporary names.

  • min_uses (int) – Minimum number of uses required before an expression is promoted to a named temporary.

Returns:

A CSEPlan containing topologically ordered assignments for reusable intermediates and the flattened scalar outputs.

Return type:

CSEPlan

gradgen.create_rust_project(function, path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64', math_library=None)[source]

Create a minimal Rust library project containing generated code.

Parameters:

  • function (Function)

  • path (str | Path)

  • config (RustBackendConfig | None)

  • crate_name (str | None)

  • function_name (str | None)

  • backend_mode (str)

  • scalar_type (str)

  • math_library (str | None)

Return type:

RustProjectResult

gradgen.create_rust_derivative_bundle(function, path, *, config=None, include_primal=True, include_jacobians=True, include_hessians=True, simplify_derivatives=None)[source]

Create a directory containing Rust crates for primal and derivatives.

Parameters:

  • function (Function)

  • path (str | Path)

  • config (RustBackendConfig | None)

  • include_primal (bool)

  • include_jacobians (bool)

  • include_hessians (bool)

  • simplify_derivatives (int | str | None)

Return type:

RustDerivativeBundleResult

gradgen.create_multi_function_rust_project(functions, path, *, config=None)[source]

Create a single Rust crate containing multiple generated kernels.

Parameters:
Return type:

RustMultiFunctionProjectResult

gradgen.derivative(expr, wrt, seed=1.0)[source]

Compute the forward derivative of a scalar expression.

Parameters:

  • expr (SX)

  • wrt (SX)

  • seed (SX | float | int)

Return type:

SX

gradgen.clear_registered_elementary_functions()[source]

Clear the custom elementary-function registry.

Return type:

None

gradgen.erf(expr)[source]

Return the symbolic error function of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.erfc(expr)[source]

Return the symbolic complementary error function of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.exp(expr)[source]

Return the symbolic exponential of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.expm1(expr)[source]

Return the symbolic exp(x) - 1 of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.floor(expr)[source]

Return the symbolic floor of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.fract(expr)[source]

Return the symbolic fractional part of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.gradient(expr, wrt, seed=1.0)[source]

Compute a reverse-mode gradient of a scalar expression.

Parameters:
Return type:

SX | SXVector

gradgen.get_registered_elementary_function(name)[source]

Return a registered elementary function by name.

Parameters:

name (str)

Return type:

RegisteredElementaryFunction

gradgen.hessian(expr, wrt)[source]

Compute a symbolic Hessian for a scalar expression.

The current return shape follows the vector-first representation:

  • scalar wrt scalar -> SX

  • scalar wrt vector -> tuple[SXVector, ...], one row per variable

Parameters:
Return type:

SX | tuple[SXVector, …]

gradgen.hypot(lhs, rhs)[source]

Return the symbolic Euclidean norm of two scalar-like values.

This is a convenience wrapper around the corresponding symbolic operator.

Parameters:

  • lhs (object) – Scalar-like left operand.

  • rhs (object) – Scalar-like right operand.

Returns:

A new SX expression representing hypot(lhs, rhs).

Return type:

SX

gradgen.if_else(if_true: object, if_false: object, condition: object) → SX[source]
gradgen.if_else(if_true: SXVector, if_false: SXVector, condition: object) → SXVector

Return a symbolic branch selected by condition.

The condition is treated numerically: zero selects if_false and any nonzero value selects if_true. Comparison expressions such as x >= y therefore work naturally as conditions.

Parameters:

  • if_true (object) – Expression returned when condition is true.

  • if_false (object) – Expression returned when condition is false.

  • condition (object) – Scalar-like selector expression.

Returns:

A scalar SX expression, or an SXVector when both branches are vectors of the same length.

Return type:

SX | SXVector

gradgen.jacobian(expr, wrt)[source]

Compute a symbolic Jacobian.

The return shape follows the current no-matrix design:

  • scalar wrt scalar -> SX

  • scalar wrt vector -> SXVector

  • vector wrt scalar -> SXVector

  • vector wrt vector -> flat row-major SXVector

Parameters:
Return type:

SX | SXVector

gradgen.jvp(expr, wrt, tangent=None)[source]

Compute a symbolic Jacobian-vector product in forward mode.

Parameters:

  • expr (SX | SXVector) – Expression to differentiate.

  • wrt (SX | SXVector) – Scalar or vector variables with respect to which expr is differentiated.

  • tangent (SX | SXVector | float | int | list[object] | tuple[object, ...] | None) – Tangent seed matching the shape of wrt. For scalar differentiation this defaults to 1.0.

Returns:

A symbolic expression with the same shape as expr containing the forward-mode derivative in the supplied direction.

Return type:

SX | SXVector

gradgen.log(expr)[source]

Return the symbolic natural logarithm of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.log1p(expr)[source]

Return the symbolic log(1 + x) of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.map_function(function, count, *, input_name=None, name=None, simplification=None)[source]

Return a staged mapped function for a unary function.

Parameters:

  • function (Function) – The symbolic function to stage. It must accept exactly one input argument.

  • count (int) – Number of repeated applications in the packed sequence.

  • input_name (str | None) – Optional name for the packed input sequence. When not provided, the name defaults to "<input_name>_seq".

  • name (str | None) – Optional name for the staged mapped function. When not provided, the generated name defaults to "<function.name>_map".

  • simplification (int | str | None) – Optional simplification effort passed through to the expanded symbolic function.

Returns:

A BatchedFunction representing the staged batched function.

Return type:

BatchedFunction

gradgen.reduce_function(function, count, *, accumulator_input_name=None, input_name=None, output_name=None, name=None, simplification=None)[source]

Return a staged left-fold reduction for a binary stage function.

Parameters:

  • function (Function) – The symbolic function to stage. It must accept exactly two inputs and produce exactly one output.

  • count (int) – Number of repeated reduction stages.

  • accumulator_input_name (str | None) – Optional name for the accumulator input. When not provided, the name defaults to "<input0>0".

  • input_name (str | None) – Optional name for the packed sequence input. When not provided, the name defaults to "<input1>_seq".

  • output_name (str | None) – Optional name for the final reduced output. When not provided, the name defaults to the stage function’s output name.

  • name (str | None) – Optional name for the staged reduction. When not provided, the generated name defaults to "<function.name>_reduce".

  • simplification (int | str | None) – Optional simplification effort passed through to the expanded symbolic function.

Returns:

A ReducedFunction describing the staged reduction.

Return type:

ReducedFunction

gradgen.matvec(matrix, x)[source]

Return the symbolic matrix-vector product for a constant matrix.

The matrix is validated as a dense rectangular numeric sequence and encoded into one matvec_component node per output row.

Parameters:

  • matrix (Sequence[Sequence[float | int]]) – Constant numeric matrix in row-major form.

  • x (object) – Symbolic vector operand with the same number of columns as matrix.

Returns:

An SXVector whose entries represent the matrix-vector product.

Raises:

ValueError – Raised when the matrix columns do not match the vector length.

Return type:

SXVector

gradgen.maximum(lhs, rhs)[source]

Return the symbolic maximum of two scalar-like expressions.

This is a convenience wrapper around the corresponding symbolic operator.

Parameters:

  • lhs (object) – First scalar-like value to compare.

  • rhs (object) – Second scalar-like value to compare.

Returns:

A new SX expression representing max(lhs, rhs).

Return type:

SX

gradgen.minimum(lhs, rhs)[source]

Return the symbolic minimum of two scalar-like expressions.

This is a convenience wrapper around the corresponding symbolic operator.

Parameters:

  • lhs (object) – First scalar-like value to compare.

  • rhs (object) – Second scalar-like value to compare.

Returns:

A new SX expression representing min(lhs, rhs).

Return type:

SX

gradgen.quadform(matrix, x, is_symmetric=True)[source]

Return a symbolic quadratic form x^   op P x for a constant matrix P.

Parameters:

  • matrix (Sequence[Sequence[float | int]]) – Constant numeric matrix in row-major form.

  • x (object) – Symbolic vector operand with the same dimension as matrix.

  • is_symmetric (bool) – When True (default), matrix is treated as symmetric: the function verifies symmetry and builds a compact equivalent form that keeps diagonal terms and doubles upper-triangular cross terms (with lower-triangular entries set to zero). When False, all entries of matrix are used.

Returns:

A scalar SX expression representing x^    op P x.

Raises:

ValueError – If matrix is not square, dimensions do not match x, or is_symmetric=True and matrix is not symmetric.

Return type:

SX

gradgen.round(expr)[source]

Return the symbolic nearest integer of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.register_elementary_function(*, name, input_dimension, parameter_dimension=0, parameter_defaults=None, eval_python=None, jacobian=None, hessian=None, hvp=None, rust_primal=None, rust_jacobian=None, rust_hvp=None, rust_hessian=None)[source]

Register a user-defined elementary function.

Parameters:

  • name (str)

  • input_dimension (int)

  • parameter_dimension (int)

  • parameter_defaults (Sequence[float | int] | None)

  • eval_python (Callable[[...], float] | None)

  • jacobian (Callable[[...], object] | None)

  • hessian (Callable[[...], object] | None)

  • hvp (Callable[[...], object] | None)

  • rust_primal (str | None)

  • rust_jacobian (str | None)

  • rust_hvp (str | None)

  • rust_hessian (str | None)

Return type:

RegisteredElementaryFunction

gradgen.generate_rust(function, *, config=None, function_name=None, backend_mode='std', scalar_type='f64', math_library=None, function_index=0, shared_helper_nodes=None, shared_helper_suppressed_custom_wrappers=None, emit_crate_header=True, function_keyword='pub fn')[source]

Generate Rust code for a symbolic function or function family.

This dispatcher inspects the supplied symbolic object and routes it to the appropriate Rust generator. It supports plain symbolic functions, composed functions, batched functions, single-shooting helpers, and the derivative variants built from those abstractions.

Parameters:

  • function (object) – The symbolic object to lower to Rust. This may be a Function, a composed function, a batched function, or one of the single-shooting helper types.

  • config (RustBackendConfig | None) – Optional Rust backend configuration. When omitted, the generator uses the default backend settings.

  • function_name (str | None) – Optional override for the generated Rust function name.

  • backend_mode (str) – Backend mode used for code generation. The default is "std".

  • scalar_type (str) – Scalar type used in the generated Rust source. The default is "f64".

  • math_library (str | None) – Optional math library name used in no_std mode.

  • function_index (int) – Index of the function within a multi-function render family.

  • shared_helper_nodes (tuple[SXNode, ...] | None) – Optional shared symbolic nodes used to keep nested helper kernels consistent.

  • shared_helper_suppressed_custom_wrappers (set[tuple[str, str]] | None) – Optional set of custom wrapper names to suppress when rendering shared helper kernels.

  • emit_crate_header (bool) – Whether to emit the crate-level header comments and attributes in the generated source.

  • function_keyword (str) – Rust keyword used for the generated function declaration. The default is "pub fn".

Returns:

A RustCodegenResult containing the rendered Rust source and its metadata.

Raises:

ValueError – If the supplied object is not supported by the Rust codegen pipeline or if the backend configuration is invalid.

Return type:

RustCodegenResult

Example

>>> from gradgen import Function, SX, generate_rust
>>> x = SX.sym("x")
>>> f = Function("square", [x], [x * x], input_names=["x"], output_names=["y"])
>>> result = generate_rust(f)
>>> result.python_name
'square'
gradgen.signum(expr)[source]

Return the symbolic sign of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.simplify(value, max_effort='basic')[source]

Simplify a symbolic expression or function.

Parameters:

  • value (Any) – Expression-like value or Function to simplify.

  • max_effort (int | str) – Maximum rewrite effort. This can be an integer pass count or one of "none", "basic", "medium", "high", or "max".

Returns:

A simplified value of the same high-level shape.

Return type:

Any

gradgen.sin(expr)[source]

Return the symbolic sine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sinh(expr)[source]

Return the symbolic hyperbolic sine of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.sqrt(expr)[source]

Return the symbolic square root of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.tan(expr)[source]

Return the symbolic tangent of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.tanh(expr)[source]

Return the symbolic hyperbolic tangent of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.trunc(expr)[source]

Return the symbolic truncation of an expression.

This is a convenience wrapper around the underlying symbolic constructors.

Parameters:

expr (object) – Scalar-like input.

Returns:

A new SX expression.

Return type:

SX

gradgen.vector(values)[source]

Create an SXVector from an iterable of scalar-like values.

Each element is coerced individually so the resulting vector is fully symbolic.

Parameters:

values (Iterable[object]) – Iterable yielding scalar-like values.

Returns:

A new SXVector containing the coerced entries.

Return type:

SXVector

gradgen.vjp(expr, wrt, cotangent=None)[source]

Compute a symbolic vector-Jacobian product in reverse mode.

Parameters:

  • expr (SX | SXVector) – Output expression being differentiated.

  • wrt (SX | SXVector) – Scalar or vector variables with respect to which expr is differentiated.

  • cotangent (SX | SXVector | float | int | list[object] | tuple[object, ...] | None) – Cotangent seed matching the shape of expr. For scalar outputs this defaults to 1.0.

Returns:

A symbolic expression with the same shape as wrt containing the reverse-mode derivative in the supplied cotangent direction.

Return type:

SX | SXVector

gradgen.zip_function(function, count, *, input_names=None, name=None, simplification=None)[source]

Return a staged batched function for a multi-input function.

This helper creates a loop-structured wrapper around function so it can be evaluated repeatedly over packed input sequences. The returned object is a BatchedFunction, which can later be expanded back into a regular symbolic Function or used for Rust code generation.

Parameters:

  • function (Function) – The symbolic function to stage. It may accept one or more inputs, and each input can be either a scalar SX or an SXVector.

  • count (int) – The number of times the function should be applied over the packed input sequences.

  • input_names (tuple[str, ...] | list[str] | None) – Optional names for the packed input sequences. When not provided, each input name defaults to "<input_name>_seq".

  • name (str | None) – Optional name for the staged batched function. When not provided, the generated name defaults to "<function.name>_zip".

  • simplification (int | str | None) – Optional simplification effort passed through to the generated symbolic function when the staged wrapper is expanded.

Returns:

A BatchedFunction describing the staged batched version of function.

Return type:

BatchedFunction

Example

>>> from gradgen import Function, SX
>>> x = SX.sym("x")
>>> y = SX.sym("y")
>>> f = Function("add", (x, y), (x + y,))
>>> batched = zip_function(f, 4)
>>> batched.name
'add_zip'
>>> batched.input_names
('x_seq', 'y_seq')
class gradgen.BatchedFunction(function, count, name, input_sequence_names, simplification=None)[source]

Bases: object

Staged wrapper that batches a function over packed inputs.

The wrapper represents a symbolic function evaluated repeatedly over packed input sequences, producing packed outputs suitable for Rust code generation or symbolic expansion.

Parameters:

  • function (Function)

  • count (int)

  • name (str)

  • input_sequence_names (tuple[str, ...])

  • simplification (int | str | None)

function

The symbolic function to stage.

Type:

gradgen.function.Function

count

Number of repeated stages in the packed sequence.

Type:

int

name

Symbolic name of the staged wrapper.

Type:

str

input_sequence_names

Packed input sequence names.

Type:

tuple[str, …]

simplification

Optional simplification effort forwarded to the expanded symbolic function.

Type:

int | str | None

function: Function
count: int
name: str
input_sequence_names: tuple[str, ...]
simplification: int | str | None
property input_names: tuple[str, ...]

Return the packed input sequence names.

property output_names: tuple[str, ...]

Return the packed output names.

property nodes

Return dependency nodes for symbolic fallback usage.

jacobian(wrt_index=0, *, name=None)[source]

Return a staged Jacobian wrapper for one packed input.

Parameters:

  • wrt_index (int) – Input index to differentiate with respect to.

  • name (str | None) – Optional symbolic name for the staged Jacobian wrapper.

Returns:

A BatchedJacobianFunction describing the staged Jacobian.

Return type:

BatchedJacobianFunction

to_function(name=None)[source]

Expand this staged batching into a regular symbolic function.

Parameters:

name (str | None) – Optional name for the expanded symbolic function. When not provided, the wrapper name is reused.

Returns:

A symbolic Function representing the expanded staged batch.

Return type:

Function

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged batched kernel.

Parameters:

  • config – Optional backend configuration object.

  • function_name (str | None) – Optional exported Rust function name.

  • backend_mode (str) – Rust backend mode to target.

  • scalar_type (str) – Rust scalar type to emit.

Returns:

The generated Rust code bundle for this staged batch.

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged batched kernel.

Parameters:

  • path (str) – Directory where the Rust crate should be created.

  • config – Optional backend configuration object.

  • crate_name (str | None) – Optional crate name override.

  • function_name (str | None) – Optional exported Rust function name.

  • backend_mode (str) – Rust backend mode to target.

  • scalar_type (str) – Rust scalar type to emit.

Returns:

A Rust project bundle rooted at path.

class gradgen.BatchedJacobianFunction(batched, wrt_index, name, simplification=None)[source]

Bases: object

Staged Jacobian wrapper for a batched function.

This object represents the Jacobian of a staged batched function without immediately expanding it back into an ordinary Function.

Parameters:

  • batched (BatchedFunction)

  • wrt_index (int)

  • name (str)

  • simplification (int | str | None)

batched

The underlying staged batched function.

Type:

gradgen.map_zip.BatchedFunction

wrt_index

The input index with respect to which the Jacobian is taken.

Type:

int

name

The symbolic name of the staged Jacobian wrapper.

Type:

str

simplification

Optional simplification effort forwarded to the expanded symbolic function.

Type:

int | str | None

batched: BatchedFunction
wrt_index: int
name: str
simplification: int | str | None
to_function(name=None)[source]

Expand this staged Jacobian into a regular symbolic function.

Parameters:

name (str | None) – Optional name for the expanded symbolic function. When not provided, the wrapper name is reused.

Returns:

A symbolic Function representing the expanded Jacobian.

Return type:

Function

property input_names: tuple[str, ...]

Return the packed input names used by the staged Jacobian.

property output_names: tuple[str, ...]

Return the generated Jacobian output names.

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged Jacobian kernel.

Parameters:

  • config – Optional backend configuration object.

  • function_name (str | None) – Optional exported Rust function name.

  • backend_mode (str) – Rust backend mode to target.

  • scalar_type (str) – Rust scalar type to emit.

Returns:

The generated Rust code bundle for this staged Jacobian.

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged Jacobian kernel.

Parameters:

  • path (str) – Directory where the Rust crate should be created.

  • config – Optional backend configuration object.

  • crate_name (str | None) – Optional crate name override.

  • function_name (str | None) – Optional exported Rust function name.

  • backend_mode (str) – Rust backend mode to target.

  • scalar_type (str) – Rust scalar type to emit.

Returns:

A Rust project bundle rooted at path.

class gradgen.ReducedFunction(function, count, name, accumulator_input_name, input_name, output_name, simplification=None)[source]

Bases: object

Staged wrapper that reduces a packed input sequence.

The wrapper represents a repeated left-fold reduction over a packed input sequence and is designed for symbolic expansion or Rust generation.

Parameters:

  • function (Function)

  • count (int)

  • name (str)

  • accumulator_input_name (str)

  • input_name (str)

  • output_name (str)

  • simplification (int | str | None)

function

The binary stage function being reduced.

Type:

gradgen.function.Function

count

Number of reduction stages.

Type:

int

name

Symbolic name of the staged reducer.

Type:

str

accumulator_input_name

Name of the accumulator input sequence.

Type:

str

input_name

Name of the packed sequence input.

Type:

str

output_name

Name of the final reduced output.

Type:

str

simplification

Optional simplification effort forwarded to the expanded symbolic function.

Type:

int | str | None

function: Function
count: int
name: str
accumulator_input_name: str
input_name: str
output_name: str
simplification: int | str | None
property input_names: tuple[str, ...]

(acc0, x_seq).

Type:

Return packed reduce input names

property output_names: tuple[str, ...]

Return the reduction output names.

property nodes

Return dependency nodes for symbolic fallback usage.

to_function(name=None)[source]

Expand this staged reduce into a regular symbolic function.

Parameters:

name (str | None) – Optional name for the expanded symbolic function. When not provided, the wrapper name is reused.

Returns:

A symbolic Function representing the expanded reduction.

Return type:

Function

generate_rust(*, config=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Generate compact Rust for the staged reduce kernel.

Parameters:

  • config – Optional backend configuration object.

  • function_name (str | None) – Optional exported Rust function name.

  • backend_mode (str) – Rust backend mode to target.

  • scalar_type (str) – Rust scalar type to emit.

Returns:

The generated Rust code bundle for this staged reduction.

create_rust_project(path, *, config=None, crate_name=None, function_name=None, backend_mode='std', scalar_type='f64')[source]

Create a Rust crate containing the staged reduce kernel.

Parameters:

  • path (str) – Directory where the Rust crate should be created.

  • config – Optional backend configuration object.

  • crate_name (str | None) – Optional crate name override.

  • function_name (str | None) – Optional exported Rust function name.

  • backend_mode (str) – Rust backend mode to target.

  • scalar_type (str) – Rust scalar type to emit.

Returns:

A Rust project bundle rooted at path.