🦀 SOC projections in Rust and other improvements

.github/pull_request_template.md

@@ -0,0 +1,33 @@
+## Main Changes
+
+Give a summary of the key changes
+
+
+## Associated Issues
+
+- Closes #1
+- Addresses #2
+
+
+## Merge Dependencies
+
+Either: 
+- This PR has no merge dependencies, or
+- Depends on: ..., or
+- Should be merged before: ...
+
+## Target versions and releases
+
+- Gradgen v0.X.X
+- Alpha release v0.X.Xa1
+
+
+## Checklist
+
+- [ ] API documentation updated
+- [ ] Appropriate testing
+- [ ] Create a demo in `demos/`, if necessary
+- [ ] Update `CHANGELOG`
+- [ ] Update webpage documentation (incl. Google Colab demo), if necessary; update `last_update` metadata
+- [ ] Bump versions in `pyproject.toml`, if necessary
+- [ ] If a new version is being released, write blog post

.github/workflows/python-tests.yml

@@ -29,7 +29,7 @@ jobs:
         run: python -m pip install --upgrade pip && python -m pip install -e . numpy sympy pytest
 
       - name: Run tests
-        run: pytest tests -v
+        run: pytest tests
 
       - name: Generate demo crates
         run: make -C demos/
@@ -51,6 +51,8 @@ jobs:
           cargo clippy --quiet --manifest-path demos/single_shooting/single_shooting_kernel/Cargo.toml -- -D warnings
           cargo clippy --quiet --manifest-path demos/single_shooting_penalty/single_shooting_penalty_kernel/Cargo.toml -- -D warnings
           cargo clippy --quiet --manifest-path demos/python_interface/foo_python/Cargo.toml -- -D warnings
+          cargo clippy --quiet --manifest-path demos/soc/soc_kernel/Cargo.toml -- -D warnings
+          cargo clippy --quiet --manifest-path demos/soc/soc_kernel_python/Cargo.toml -- -D warnings
 
       - name: Run demo runners
         run: |

AGENTS.md

@@ -111,7 +111,7 @@ def zip_function(
 
 ### Docusaurus
 
-- All docusaurus files must be `.mdx` and not `.md`. Never rename `.mdx` files to `.md`.
+- All docusaurus files, with the possible exception of blog posts, must be `.mdx` and not `.md`. Never rename `.mdx` files to `.md`.
 
 ## What Agents Must Never Do
 

CHANGELOG.md

@@ -6,6 +6,58 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/)
 and this project adheres to [Semantic Versioning](http://semver.org/).
 
 
+## 0.5.1 - 06-05-2026
+
+### Added
+
+- Added `SquaredDistanceToSet.second_order_cone(...)` for Rust-backed
+  half-squared distance and projection generation for second-order cones
+  of the form `{x = (y, t) : ||y||_2 <= alpha t}`.
+- Added a `demos/soc` demo that generates a Rust-Python interface for a
+  second-order cone distance kernel and prints the generated function
+  outputs from Python.
+- Added SymPy parity integration tests covering scalar derivatives,
+  gradients, Jacobians, and Hessians generated by gradgen.
+
+### Changed
+
+- Removed the `with_prefer_direct_output_sinks` Rust backend config toggle and
+  made additive output accumulation the default lowering style for final
+  outputs. Generated kernels now write readable incremental updates directly
+  into the output slices whenever possible.
+- Emitted `copy_from_slice` for contiguous vector outputs backed by workspace
+  slices, replacing element-by-element final copy assignments with idiomatic
+  slice copies where possible.
+- Optimized the generated Rust helpers for
+  `SquaredDistanceToSet.second_order_cone(...)` so the projection defers
+  `sqrt(...)` to the nontrivial branch, and the primal helper computes the
+  half-squared distance directly without first materializing the projection.
+- Fixed Rust-backed custom snippets used by
+  `SquaredDistanceToSet.second_order_cone(...)` so generated `no_std`
+  crates use the configured math library for `sqrt(...)` instead of calling
+  an unavailable inherent floating-point method.
+- Optimized generated Rust gradients for projection-backed custom distances
+  so trivial `+ 0` and `* 1` wrappers no longer force workspace code, direct
+  Jacobian helper calls can be emitted again, and synthesized Jacobian
+  helpers now use loops instead of fully unrolled per-component assignments.
+- Fixed workspace liveness tracking in Rust codegen so shared intermediates
+  are not overwritten before their final consumers, restoring correct
+  generated derivatives and other workspace-heavy kernels.
+
+### Fixed
+
+- Fixed Rust rendering of squared expressions so non-atomic bases are
+  parenthesized before applying the square fast path. This restores correct
+  code generation for terms such as `(x - p)^2` in simplified single-shooting
+  costs and similar expressions.
+- Optimized `SquaredDistanceToSet.euclidean_ball(...)` Rust code generation by
+  routing it through compact Rust helper snippets and Python callbacks,
+  avoiding scalarized symbolic expansions such as `x[i] - 0` chains and
+  enabling loop-based primal and Jacobian kernels.
+- Optimized `SquaredDistanceToSet.infinity_ball(...)` Rust code generation in
+  the same way, including a zero-center fast path and compact loop-based
+  primal, projection, and Jacobian helpers.
+
 
 ## 0.5.0 - 01-05-2026
 

demos/Makefile

@@ -1,11 +1,11 @@
-.PHONY: all codegen custom_function squared_distance_to_set squared_distance_to_set_rust_only multi_function vjp map_zip zip_3 reduce_map composed_function composed_chain single_shooting single_shooting_penalty python_interface help
+.PHONY: all codegen custom_function squared_distance_to_set squared_distance_to_set_rust_only multi_function vjp map_zip zip_3 reduce_map composed_function composed_chain single_shooting single_shooting_penalty python_interface soc help
 
 PYTHON ?= python
 REPEATS ?= 5
 HORIZON ?= 5
 COUNT ?= 3
 
-all: codegen custom_function squared_distance_to_set squared_distance_to_set_rust_only multi_function vjp map_zip zip_3 reduce_map composed_function composed_chain single_shooting single_shooting_penalty python_interface
+all: codegen custom_function squared_distance_to_set squared_distance_to_set_rust_only multi_function vjp map_zip zip_3 reduce_map composed_function composed_chain single_shooting single_shooting_penalty python_interface soc
 
 help:
 	@echo "Available targets:"
@@ -24,6 +24,7 @@ help:
 	@echo "  make single_shooting HORIZON=$(HORIZON)"
 	@echo "  make single_shooting_penalty HORIZON=$(HORIZON)"
 	@echo "  make python_interface"
+	@echo "  make soc"
 
 codegen:
 	$(PYTHON) codegen/main.py
@@ -66,3 +67,6 @@ single_shooting_penalty:
 
 python_interface:
 	$(PYTHON) python_interface/main.py
+
+soc:
+	$(PYTHON) soc/main.py

demos/README.md

@@ -19,6 +19,7 @@ This directory contains runnable end-to-end examples showing how to use
 - [single_shooting](./single_shooting/README.md): loop-based deterministic single-shooting OCP code generation
 - [single_shooting_penalty](./single_shooting_penalty/README.md): single-shooting OCP code generation with vector residual penalties
 - [python_interface](./python_interface/README.md): generate a Rust crate plus a separate PyO3 wrapper imported from Python
+- [soc](./soc/README.md): second-order cone distance generation with an importable Python wrapper
 
 ## Running all demos
 

demos/codegen/codegen_kernel/README.md

@@ -14,12 +14,12 @@ cargo build
 
 ## Generated Functions
 
-- `codegen_kernel_energy_f`: workspace `2`, inputs `(3, 1)`, outputs `(1,)`
-- `codegen_kernel_energy_jf_x`: workspace `3`, inputs `(3, 1)`, outputs `(3,)`
-- `codegen_kernel_energy_jf_u`: workspace `1`, inputs `(3, 1)`, outputs `(1,)`
-- `codegen_kernel_coupling_f`: workspace `2`, inputs `(3, 1)`, outputs `(1,)`
+- `codegen_kernel_energy_f`: workspace `0`, inputs `(3, 1)`, outputs `(1,)`
+- `codegen_kernel_energy_jf_x`: workspace `0`, inputs `(3, 1)`, outputs `(3,)`
+- `codegen_kernel_energy_jf_u`: workspace `0`, inputs `(3, 1)`, outputs `(1,)`
+- `codegen_kernel_coupling_f`: workspace `0`, inputs `(3, 1)`, outputs `(1,)`
 - `codegen_kernel_coupling_jf_x`: workspace `0`, inputs `(3, 1)`, outputs `(3,)`
-- `codegen_kernel_coupling_jf_u`: workspace `1`, inputs `(3, 1)`, outputs `(1,)`
+- `codegen_kernel_coupling_jf_u`: workspace `0`, inputs `(3, 1)`, outputs `(1,)`
 
 The generated ABI uses input slices, output slices, and a mutable workspace slice.
 Each generated function also includes metadata helpers for workspace, input, and output sizes unless disabled in the backend configuration.

demos/codegen/codegen_kernel/metadata.json

@@ -1,11 +1,11 @@
 {
   "crate_name": "codegen_kernel",
-  "created_at": "2026-05-01T08:32:08.672636Z",
-  "gradgen_version": "0.4.2a2",
+  "created_at": "2026-05-05T23:14:18.278994Z",
+  "gradgen_version": "0.5.1a1",
   "functions": [
     {
       "function_name": "codegen_kernel_energy_f",
-      "workspace_size": 2,
+      "workspace_size": 0,
       "input_names": [
         "x",
         "u"
@@ -23,7 +23,7 @@
     },
     {
       "function_name": "codegen_kernel_energy_jf_x",
-      "workspace_size": 3,
+      "workspace_size": 0,
       "input_names": [
         "x",
         "u"
@@ -41,7 +41,7 @@
     },
     {
       "function_name": "codegen_kernel_energy_jf_u",
-      "workspace_size": 1,
+      "workspace_size": 0,
       "input_names": [
         "x",
         "u"
@@ -59,7 +59,7 @@
     },
     {
       "function_name": "codegen_kernel_coupling_f",
-      "workspace_size": 2,
+      "workspace_size": 0,
       "input_names": [
         "x",
         "u"
@@ -95,7 +95,7 @@
     },
     {
       "function_name": "codegen_kernel_coupling_jf_u",
-      "workspace_size": 1,
+      "workspace_size": 0,
       "input_names": [
         "x",
         "u"