zip_3

Zip-3 Function Demo

This demo shows zip_function(...) with a stage function that has three inputs, where the first input is a vector.

The stage function $h:\mathbb{R}^2 \times \mathbb{R} \times \mathbb{R} \to \mathbb{R}$ given by

$$h(a, b, c) = a_1 b + a_2 + \sin(c),$$

with stage input types $a \in \mathbb{R}^2$, $b \in \mathbb{R}$, and $c \in \mathbb{R}$.

For an integer $N$ (count = N), the batched kernel computes

$$ ((a_1,\dots,a_N),\ (b_1,\dots,b_N),\ (c_1,\dots,c_N)) \mapsto \bigl(h(a_1,b_1,c_1),\dots,h(a_N, b_N, c_N)\bigr). $$

The demo also builds a staged Jacobian source for the batched function with respect to the sequence of a's, that is, $(a_1, \ldots, a_N)$ (wrt_index=0).

The code

In main.py, we define symbolic variables a, b, c and the 3-input stage function h:

# Symbols
# a: R^2, b: R, c: R
a = SXVector.sym("a", 2)
b = SX.sym("b")
c = SX.sym("c")

# Function h: R^2 x R x R --> R
h = Function(
    "h",
    [a, b, c],
    [a[0] * b + a[1] + sin(c)],
    input_names=["a", "b", "c"],
    output_names=["y"],
)

Then we define the batched function:

batched = zip_function(
    h,
    count,
    input_names=("a_seq", "b_seq", "c_seq"),
    name="zip3",
)

To evaluate it as a regular Function:

batched_fn = batched.to_function()

To build a Jacobian wrt a_seq:

batched_jac_a = batched.jacobian(wrt_index=0)
batched_jac_a_fn = batched_jac_a.to_function()

Code generation

Code generation is as with regular functions...

project = (
    CodeGenerationBuilder()
    .with_backend_config(
        RustBackendConfig()
        .with_crate_name("zip_3_kernel")
        .with_backend_mode("no_std")
        .with_scalar_type("f64")
    )
    .for_function(batched)
        .add_primal()
        .add_jacobian()
        .with_simplification("medium")
        .done()
    .build("./my_crates")
)

Files

• main.py: defines a 3-input batched staged function and generates Rust
• zip_3_kernel/: generated by running main.py
• runner/: small Rust binary crate using the generated kernels

Run

From repository root, with an activated virtual environment:

python demos/zip_3/main.py --count 3

Then run the Rust consumer:

cd demos/zip_3/runner
cargo run