# JAX-based Rosenbrock function minimization

**URL:** <https://si-tesseract.discourse.group/t/jax-based-rosenbrock-function-minimization/48>\
**Category:** Expert showcase\
**Tags:** optimization, v0-7-x\
**Created:** [February 26, 2025, 2:11pm UTC](https://si-tesseract.discourse.group/t/jax-based-rosenbrock-function-minimization/48 "2025-02-26T14:11:33Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![Niklas](https://avatars.discourse-cdn.com/v4/letter/n/34f0e0/32.png) [@Niklas](https://si-tesseract.discourse.group/u/Niklas)\
**Post date:** [February 26, 2025, 2:11pm UTC](https://si-tesseract.discourse.group/t/jax-based-rosenbrock-function-minimization/48/1 "2025-02-26T14:11:33Z")

</div>

This example demonstrates how to:

1. Create a Tesseract that wraps the [Rosenbrock function](https://en.wikipedia.org/wiki/Rosenbrock_function), a non-convex function that is commonly used to test optimization algorithms.

2. Use the tesseract with Scipy’s `minimize` routine to find the minimum of the Rosenbrock function.

The whole Tesseract can be downloaded here: [univariate.zip](https://si-tesseract.discourse.group/uploads/short-url/adYR074KBpiDCITjVJrdxdeCv2D.zip) (3.9 KB)

![Optimization trajectory on the Rosenbrock tesseract.](https://canada1.discourse-cdn.com/flex009/uploads/si_tesseract/original/1X/85959572de67e17505a3e3aef2e9da66da0ec0e4.gif)

> 🚀 For many more examples, see the [`examples/`](https://github.com/pasteurlabs/tesseract-core/tree/main/examples) directory in the Tesseract repository

## Tesseract definition

In this section we walk you through each step to re-create the example in [`examples/unit_tesseract/univariate/`](https://github.com/pasteurlabs/tesseract-core/tree/main/examples/unit_tesseracts/univariate).

The easiest way to set up a new tesseract is with the `tesseract init` command, which will create  
three files:

```bash
$ tesseract init --name univariate --target-dir your/path/to/univariate
$ ls your/path/to/univariate
tesseract_config.yaml
tesseract_requirements.txt
tesseract_api.py

```

### `tesseract_config.yaml`

The config file of the tesseract. Here you can specify the docker `base_image`, extra (non-python) packages you might want to install, and much more. We modify the config that was generated by `tesseract init`:

```yml
name: "univariate"
version: "0.1.0"
description: |
  Unit tesseract that evaluates the rosenbrock function and derivatives.

```

### `tesseract_requirements.txt`

Specifies all python dependecies of your tesseract.

### `tesseract_api.py`

The actual implementation of the tesseract. In the first few lines we define the `rosenbrock` function itself, then we specify the `InputSchema` and `OutputSchema`, which essentially mimic the type signature of `rosenbrock`.

```auto
from pydantic import BaseModel, Field
from tesseract_core.runtime import Differentiable, Float64, ShapeDType

def rosenbrock(x: float, y: float, a: float = 1.0, b: float = 100.0) -> float:
    return (a - x) **2 + b * (y - x**2) ** 2

class InputSchema(BaseModel):
    x: Differentiable[Float64] = Field(description="Scalar value x.", default=0.0)
    y: Differentiable[Float64] = Field(description="Scalar value y.", default=0.0)
    a: Float64 = Field(description="Scalar parameter a.", default=1.0)
    b: Float64 = Field(description="Scalar parameter b.", default=100.0)

class OutputSchema(BaseModel):
    result: Differentiable[Float64] = Field(
        description="Result of Rosenbrock function evaluation."
    )

def apply(inputs: InputSchema) -> OutputSchema:
    """Evaluates the Rosenbrock function given input values and parameters."""
    result = rosenbrock(inputs.x, inputs.y, a=inputs.a, b=inputs.b)
    return OutputSchema(result=result)

```

At this point you could already build the tesseract and evaluate the Rosenbrock function with

```bash
$ tesseract build your/path/to/univariate
$ tesseract run univariate apply '{"inputs": {"x":1, "y":1}}'
{
  "result": {
    "object_type": "array",
    "shape": [],
    "dtype": "float64",
    "data": {
      "buffer": 0.0,
      "encoding": "json"
    }
  }
}

```

The `result` contains a json encoded singleton array (with shape `[]`).

> 🚀 For the implementation of the AD endpoints, check the full [`tesseract_api.py`](https://github.com/pasteurlabs/tesseract-core/blob/main/examples/univariate/tesseract_api.py). It contains naive implementations of `jacobian`, `jvp`, and `vjp` to illustrate how to define your own AD endpoints. In case your `apply` endpoint is implemented in JAX, you can make use of the [Tesseract JAX template](https://si-tesseract.discourse.group/t/jax-auto-diff-templates-gaussian-radial-basis-function-fitting/51) (via `tesseract init --recipe jax ...`) to auto-generate your AD endpoints.

## Tesseract usage via the Python API

To use a built tesseract, you can serve it by name, which will make it available at `http://localhost:<PORT>`.

```bash
$ tesseract serve univariate --port 58354
 [i] Waiting for Tesseract containers to start ...
 [i] Container ID: 45c84bca8caa2d39e27fc1b01129c17f51ad19f781ede07b9275623fb4da65d2
 [i] Name: tesseract-flivhk7muce0-sha256-glfafackzws6-1
 [i] Entrypoint: ['tesseract-runtime', 'serve']
 [i] View Tesseract: http://localhost:58354/docs
 [i] Docker Compose Project ID, use it with 'tesseract teardown' command: tesseract-flivhk7muce0
{"project_id": "tesseract-flivhk7muce0", "containers": [{"name": "tesseract-flivhk7muce0-sha256-glfafackzws6-1", "port": "58354"}]}%

```

The served tesseract can be conveniently used via the Python API (the code below uses the [`animate`](https://github.com/pasteurlabs/tesseract-core/blob/main/examples/univariate/animate.py) module from the [`univariate`](https://github.com/pasteurlabs/tesseract-core/tree/main/examples/univariate) example).

```python
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import minimize

from tesseract_core import Tesseract

from animate import make_animation

# Instantiate python api to the tesseract (read port from logs above)
tesseract = Tesseract(url="http://localhost:58354")

def rosenbrock(x: np.ndarray) -> float:
    """Wrap tesseract.apply to adhere to scipy's minimize interface."""
    output = tesseract.apply({"x": x[0], "y": x[1]})
    return output["result"].item()

def rosenbrock_gradient(x: np.ndarray) -> np.ndarray:
    """Wrap tesseract.jacobian to adhere to scipy's minimize interface."""
    output = tesseract.jacobian(
        {"x": x[0], "y": x[1]}, jac_inputs=["x", "y"], jac_outputs=["result"]
    )
    return np.array([output["result"]["x"], output["result"]["y"]])

# Initial guess for the variables
x0 = np.array([-0.5, 1.0])
trajectory = [x0]

# kick off scipy's optimization using tesseract apply/jacobian
result = minimize(
    rosenbrock,
    x0,
    method="BFGS",
    jac=rosenbrock_gradient,
    options={"disp": True},
    callback=lambda xs: trajectory.append(xs.tolist()),
)

anim = make_animation(*list(zip(*trajectory)))
# anim.save("rosenbrock_optimization.gif", writer="pillow", fps=2, dpi=150)
plt.show()

```