v2 WIP

.github/workflows/ci_v2.yaml

@@ -0,0 +1,92 @@
+name: L4CasADi v2
+
+on:
+  push:
+    branches: [ v2 ]
+
+jobs:
+  lint:
+    name: Lint
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        ref: 'v2'
+    - name: Run mypy
+      run: |
+        pip install mypy
+        mypy . --ignore-missing-imports --exclude examples
+    - name: Run flake8
+      run: |
+        pip install flake8
+        # stop the build if there are Python syntax errors or undefined names
+        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+
+  tests:
+    runs-on: ${{ matrix.runs-on }}
+    needs: [ lint ]
+    timeout-minutes: 60
+    strategy:
+      fail-fast: false
+      matrix:
+        runs-on: [ubuntu-latest, ubuntu-20.04, macos-latest]
+
+    name: Tests on ${{ matrix.runs-on }}
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v3
+      with:
+        ref: 'v2'
+        fetch-depth: 0
+
+    - name: Install Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '>=3.9 <3.12'
+
+    - name: Install L4CasADi
+      run: |
+        python -m pip install --upgrade pip
+        pip install torch --index-url https://download.pytorch.org/whl/cpu  # Ensure CPU torch version
+        pip install -r requirements_build.txt
+        pip install . -v --no-build-isolation 
+
+    - name: Test with pytest
+      working-directory: ./tests
+      run: |
+        pip install pytest
+        pytest .
+
+  test-on-aarch:
+    runs-on: ubuntu-latest
+    needs: [ lint ]
+    timeout-minutes: 60
+
+    name: Tests on aarch64
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+        with:
+          ref: 'v2'
+          fetch-depth: 0
+      - uses: uraimo/run-on-arch-action@v2
+        name: Install and Test
+        with:
+          arch: aarch64
+          distro: ubuntu20.04
+          install: |
+            apt-get update
+            apt-get install -y --no-install-recommends python3.9 python3-pip python-is-python3
+            pip install -U pip
+            apt-get install -y build-essential
+
+          run: |
+            python -m pip install --upgrade pip
+            pip install torch --index-url https://download.pytorch.org/whl/cpu  # Ensure CPU torch version
+            pip install -r requirements_build.txt
+            pip install . -v --no-build-isolation 
+            # pip install pytest
+            # pytest .

README.md

@@ -1,5 +1,6 @@
 [![PyPI version](https://badge.fury.io/py/l4casadi.svg)](https://badge.fury.io/py/l4casadi)
 ![L4CasADi CI](https://github.com/Tim-Salzmann/l4casadi/actions/workflows/ci.yaml/badge.svg)
+![L4CasADi v2](https://github.com/Tim-Salzmann/l4casadi/actions/workflows/ci_v2.yaml/badge.svg)
 ![Downloads](https://img.shields.io/pypi/dm/l4casadi.svg)
 
 ---
@@ -23,6 +24,22 @@ arXiv: [Learning for CasADi: Data-driven Models in Numerical Optimization](https
 
 Talk: [Youtube](https://youtu.be/UYdkRnGr8eM?si=KEPcFEL9b7Vk2juI&t=3348)
 
+## L4CasADi v2 Changelog
+After feedback from first use-cases L4CasADi v2 is designed with efficiency and simplicity in mind. This leads to a few breaking changes.
+
+Breaking changes are
+- L4CasADi will not change the shape of an input anymore. The tensor forwarded to the PyTorch model will resemble the exact
+dimension of the input variable by CasADi. You are responsible to make sure that the PyTorch model handles a 
+**two-dimensional** input matrix!
+- L4CasADi v2 can leverage PyTorch's batching capabilities for increased efficiency. When passing `batched=True`,
+L4CasADi will understand the **first** input dimension as batch dimension. Thus, first and second-order derivatives
+across elements of this dimension are assumed to be **sparse-zero**. To make use of this, instead of having multiple calls to a L4CasADi function in
+your CasADi program, batch all inputs together and have a single L4CasADi call. An example of this can be seen when
+comparing the [non-batched NeRF example](examples/nerf_trajectory_optimization/nerf_trajectory_optimization.py) with the
+[batched NeRF example](examples/nerf_trajectory_optimization/nerf_trajectory_optimization_batched.py) which is faster by
+a factor of 5-10x.
+
+
 ## Table of Content
 - [Projects using L4CasADi](#projects-using-l4casadi)
 - [Installation](#installation)
@@ -202,14 +219,6 @@ https://github.com/Tim-Salzmann/l4casadi/blob/421de6ef408267eed0fd2519248b2152b6
 
 ## FYIs
 
-### Batch Dimension
-
-If your PyTorch model expects a batch dimension as first dimension (which most models do) you should pass
-`model_expects_batch_dim=True` to the `L4CasADi` constructor. The `MX` input to the L4CasADi component is then expected
-to be a vector of shape `[X, 1]`. L4CasADi will add a batch dimension of `1` automatically such that the input to the
-underlying PyTorch model is of shape `[1, X]`.
-
----
 
 ### Warm Up
 

examples/naive/readme.py

@@ -3,15 +3,15 @@
 
 
 naive_mlp = l4c.naive.MultiLayerPerceptron(2, 128, 1, 2, 'Tanh')
-l4c_model = l4c.L4CasADi(naive_mlp, model_expects_batch_dim=True)
+l4c_model = l4c.L4CasADi(naive_mlp)
 
-x_sym = cs.MX.sym('x', 2, 1)
+x_sym = cs.MX.sym('x', 1, 2)
 y_sym = l4c_model(x_sym)
 f = cs.Function('y', [x_sym], [y_sym])
 df = cs.Function('dy', [x_sym], [cs.jacobian(y_sym, x_sym)])
 ddf = cs.Function('ddy', [x_sym], [cs.hessian(y_sym, x_sym)[0]])
 
-x = cs.DM([[0.], [2.]])
+x = cs.DM([[0., 2.]])
 print(l4c_model(x))
 print(f(x))
 print(df(x))

examples/nerf_trajectory_optimization/nerf_trajectory_optimization.py

@@ -86,7 +86,7 @@ def trajectory_generator_solver(n, n_eval, L, warmup, threshold):
             f += cs.sum2(sk**2)
 
             # While having a maximum density (1.) of the NeRF as constraint.
-            lk = L(pk.T)
+            lk = L(pk)
             g = cs.horzcat(g, lk)
             lbg = cs.horzcat(lbg, cs.DM([-10e32]).T)
             ubg = cs.horzcat(ubg, cs.DM([threshold]).T)