JAX-based Rosenbrock function minimization

Showcase Expert showcase
,
1

This example demonstrates how to:

  1. Create a Tesseract that wraps the 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 (3.9 KB)

Optimization trajectory on the Rosenbrock tesseract.

:rocket: For many more examples, see the 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/.

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

$ 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:

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.

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

$ 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 []).

:rocket: For the implementation of the AD endpoints, check the full 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 (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>.

$ 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 module from the univariate example).

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()
1 Like