{"slug": "a-tour-of-differentiable-rasterization", "title": "A Tour of Differentiable Rasterization", "summary": "Sasha Rush's notebook introduces differentiable rasterization, a method to make the conversion of vector graphics to pixels differentiable for gradient-based optimization in deep learning frameworks like Jax and PyTorch. The post demonstrates how to represent vector graphics as affine transforms and use autodiff to adjust parameters, such as ellipses, to match a target image. It targets ML practitioners and is built with Chalk, developed with Dan Oneață.", "body_md": "*By Sasha Rush - Notebook*\n\n*Built with Chalk developed with Dan Oneață*\n\n*Based on*\n\nGiven a program that produces a vector representation of an image (think SVG), *rasterization* turns it into a pixel representation (think PNG). If we have a target image, we can see how close we got.\n\n\\[\\begin{align*} \\text{vector} &= \\text{program}(x) \\\\ \\text{im} &= \\text{rasterizer}(\\text{vector}) \\\\ {\\cal L}(\\text{im}, \\text{target}) &= \\| \\text{im} - \\text{target} \\| \\\\ \\end{align*}\\]\n\nUnfortunately this process is not inherently differentiable. Moreover, it doesn’t easily run in modern deep learning frameworks like Jax or PyTorch. The goal of this post is to walk through this calculation.\n\n\\[\\text{Goal:} \\frac{d {\\cal L}}{dx} = \\frac{d {\\cal L}}{d\\ \\text{im}} \\frac{d\\ {\\text{im}}}{d\\ \\text{vector}} \\frac{d\\ \\text{vector}}{dx}\\]\n\nTo make this all more tangible, say we have a target image and a program that can draw ellipses based on their locations, size, and rotation. Differentiable rasterization would let us use gradient descent to adjust the parameters \\(x\\) to match the image.\n\nThis blog is not really about computer graphics (I don’t know much about computer graphics). The goal is to explore *differentiable programming* in realistic settings. If autodiff + vectorization is the future, then it is important to be able to write hard programs in a differentiable style (beyond just another Transformer).\n\nThe blog is in 5 parts. It assumes no graphics knowledge, but does use a lot of Jax/NumPy tricks. The target audience is someone with a lot of ML experience who wants to branch out into some more complex uses of derivatives.\n\nThis section is about vector graphics and linear algebra. We want to be able to represent a vector graphic as a function of parameters using matrices.\n\n\\[\\begin{align*} \\text{vector} &= \\text{program}(x) \\\\ \\end{align*} \\]\n\nTo arrive here, we need to convince ourselves that this process is fully differentiable and expressible in NumPy. This will also allow us to introduce the basic primitives of vector graphics.\n\nWe start with a *transform*. This is an affine function of a 2D point. Given a point \\((x, y)\\), we produce a point \\((x', y')\\) where \\(x' = a x + b y + c\\) and \\(y' = d x + e y + f\\). This will allow us to mathematically represent intuitive notions of scaling, rotation, and translation.\n\nWe can represent these affine functions as a 3x3 matrices. \\[\\begin{align*} \\begin{bmatrix} x' \\\\ y' \\\\ 1 \\end{bmatrix} &= \\begin{bmatrix} a & b & c \\\\ d & e & f \\\\ 0 & 0 & 1 \\end{bmatrix} \\begin{bmatrix} x \\\\ y \\\\ 1 \\end{bmatrix} \\\\ \\end{align*}\\] The key trick being to represent 2D points in 3D with a fixed 1 in the last position. We can the apply transforms to points, lines, shapes, or other transforms etc. As these are standard matrix operations, they naturally have meaningful derivatives, which autodiff can provide for us.\n\nWe start by applying them to a simple line.\n\n```\nline = arc_seg(V2(1, 1), 1e-3).stroke()\n```\n\nTranslations have an identity scaling and use the last column for the offet. We show the current transform matrix on the right side.\n\n``` python\ndef translate(t: float):\n    \"t is a float between 0 and 1\"\n    affine = tx.translation(V2(0, t))\n    return line.apply_transform(affine), show_affine(affine)\n\nanimate(translate)\n```\n\nScaling transforms use the top left diagonal of the matrix.\n\n``` python\ndef scale(t):\n    affine = tx.scale(t)\n    return line.apply_transform(affine), show_affine(affine)\n\nanimate(scale)\n```\n\nRotation transforms uses the off-diagonal.\n\n``` python\ndef rotate(t):\n    affine = tx.rotation(t * 2 * np.pi)\n    return line.apply_transform(affine), show_affine(affine)\n\nanimate(rotate)\n```\n\nIn our little library we use `+`\n\nto represent composition of vector diagrams. Here two lines are combined in a single image. The first is rotated and the second is rotated in the opposite direction using the matrix inverse of the transform matrix.\n\n``` python\ndef cat(t):\n    affine = tx.rotation(t * 2 * np.pi)\n    affine2 = tx.inv(affine)\n    return line.apply_transform(affine) + line.apply_transform(affine2), show_affine(\n        affine\n    ) | show_affine(affine2)\n\nanimate(cat)\n```\n\nColors are also represented mathematically. We use simple RGB vectors.\n\n\\[ \\text{color} = \\begin{bmatrix} r \\\\ g \\\\ b \\end{bmatrix}\\]\n\n``` python\ndef color_line(t):\n    # Helper function for standard colors. \n    red = to_color(\"red\")\n    blue = to_color(\"blue\")\n    color = t * red + (1 - t) * blue\n    return line.line_color(color), show_color(color)\n\nanimate(color_line)\n```\n\nFor simplicity, the only function we provide is an *arc*. Arcs are created by providing a vector offset and the bend height. Note that lines are just arcs with a bend near 0.\n\n``` python\ndef arc(t):\n    # End position of the arc.\n    v = V2(1, 0)\n    # Create an arc from (0, 0) to v with bend t.\n    a_s = arc_seg(v, t)\n    # Stroke turns an arc into a vector diagram. \n    return a_s.stroke()\n\nanimate(arc)\n```\n\nInternally arcs are circles clipped between two angles. We can do this calculation with some (differentiable) [high-school trigonometry](https://observablehq.com/@sarah37/arcs-between-two-points).\n\nHere is what these arcs look like by composing the arc and the internal circle representation.\n\n``` python\ndef arc(t):\n    v = V2(1, 0)\n    a_s = arc_seg(v, t)\n    return a_s.stroke().line_width(2) + show_arc(a_s).stroke().line_width(0.1)\n\nanimate(arc, lw=False)\n```\n\nWe can *chain* arcs together before stroking them to vector diagrams. We agin use `+`\n\noperator for chaining before closing the path.\n\n(Note that `+`\n\nis used both for chaining and composition. This is because both of [monoids](https://core.ac.uk/download/pdf/76383233.pdf) and so they share internal structure. This can be point of confusion though.)\n\nFilled paths made with chained arcs will be the core element we work with.\n\n``` python\ndef closed(t):\n    v = V2(1, 0)\n    return (arc_seg(v, t) + arc_seg(-v, -0.2)).close().stroke()\n\nanimate(closed)\n```\n\nWe conclude with a circle function.\n\n``` python\ndef circle_at(p, r):\n    \"Draw a circle at `p` with radius `r`\"\n    # Draw a circle with radius 1 centered at (0, 1)\n    circle = (arc_seg(V2(0, 2), 1) + arc_seg(V2(0, -2), 1)).close().stroke()\n    # Move the circle to p and scale to radius r\n    affine = tx.translation(p) @ tx.scale(r) @ tx.translation(V2(0, -1))\n    return circle.apply_transform(affine)\n\ndef draw_circle(t):\n    return circle_at(t, t)\n\nanimate(draw_circle)\n```\n\nThis section is about Jax, a library for differentiable programming. Jax is built on top of NumPy so it works directly with NumPy arrays you probably are used to. Basically python dataclass that is made up of NumPy arrays can be utilized withing the Jax framework. The key thing to note is that our vector diagrams are [Jax PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html).\n\nThe downside of Jax is that we will not be able to use a lot of standard python functions, and will have to write everything in a strict functional style without `if's`\n\n. This will make later things harder.\n\nBut there are upsides. In this section we will discuss the use of *vmap*, which allows us to automatically vectorize functions, and *grad* which will allow us to compute derivatives are arbitrary functions. # Let’s start with [vmap](https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html). This function will allow us to easily create complex vector diagrams, by describing many different paths at once. For example here we are using `vmap`\n\nto create 10 circles.\n\n(Note from here on out we use simpler methods like `scale_x`\n\ninstead of manually creating affine matrices. Hopefully they are self-explanatory)\n\n``` python\ndef draw(t):\n    @jax.vmap\n    def multi(i):\n        \"Draw a circle with radius using `i`.\"\n        return circle_at(P2(0, 0), (0.05 * i + t)).scale_x(1.2).fill_opacity(0)\n\n    # Create 9 circles with i from 1 to 9.\n    return multi(np.arange(1, 10))\n\nanimate(draw)\n```\n\nYou might have noticed the code above for simple functions is kind of slow. Jax transforms take a little longer to startup, but they automatically let us scale to some very complex figures. Here we are creating 100 circles.\n\n``` python\ndef draw(t):\n    red = to_color(\"red\")\n    green = to_color(\"green\")\n\n    @jax.vmap\n    def multi(i):\n        i_s = i / 100\n        s = (0.05 * i) / 2\n        rot = tx.rotation(i * t + i)\n        return (\n            rectangle(s, s)\n            .fill_color(i_s * red + (1 - i_s) * green)\n            .line_width(0)\n            .apply_transform(rot)\n        )\n\n    return multi(np.arange(100, 1, -1))\n\nanimate(draw, lw=False, steps=200, rate=5)\n```\n\nWe can use vmap on multiple arguments simulataneously or even entire trees of parameters. Here is one more fun example taking multiple arguments.\n\n``` python\nT = 100\n\ndef multi(t):\n    r, g = to_color(\"red\"), to_color(\"green\")\n\n    @jax.vmap\n    def draw(x, y):\n        t2 = (x + y + t) % 1\n        t2 = np.maximum(t2, 1 - t2)\n        return (\n            circle_at(P2(x, y), t2 / 10 - 0.2)\n            .fill_color(t2 / 2 * r + 2 * (1 - t2) * g)\n            .translate(-0.5, -0.5)\n            .scale(2)\n            .line_width(0)\n        )\n\n    return draw((np.arange(T) // 10) / 10, (np.arange(T) % 10) / 10)\n\nanimate(jax.jit(multi), steps=10, lw=False)\n```\n\nThe other core benefit of having Jax is to be able to differentiate functions. Since we are able to take derivatives of all our functions we can use gradient descent to optimize for arbitrary properties.\n\nWe’ll introduce a function `opt`\n\nthat uses [Adam](https://optax.readthedocs.io/en/latest/api/optimizers.html#optax.adam) to optimize and draw an arbitrary function. Our optimization uses unconstrained optimization with\n\na starting point and a function `f`\n\n, and draws the intermediate vector diagrams.\n\n\\[\\arg\\min_{x} f(x)\\]\n\nThere are more interesting optimizers we could apply here, but we’ll keep it simple for now.\n\n``` python\ndef L2(x):\n    return (x * x).sum(-1)\n\ndef find_point(target, x):\n    return L2(target - x).sum(), (\n        grid\n        + circle_at(target, 0.1).fill_color(\"red\")\n        + circle_at(x, 0.1).fill_color(\"blue\")\n    ).layout(500)\n\n# Use partial to specify the target. Optimize over x.\nopt(P2(0, 0), partial(find_point, P2(0.5, 0.5)), steps=100, rate=0.1)\n```\n\nThe two aspects work well together. A [force-directed graph](https://en.wikipedia.org/wiki/Force-directed_graph_drawing) is a way to layout a graph in a visually pleasing way in 2D space. The approach for physical forces over the nodes and edges.\n\nIn this function we use three forces. A spring along edges, repulsion between all nodes, and a weak gravity towards the center of the graph.\n\n``` python\n@jax.jit\ndef force_directed(x, edges, a=1 / 20, b=1 / 10, c=50, spring=0.04):\n    size = x.shape[0]\n    nodes = np.arange(size)\n    # Colors for nodes\n    color = np.stack([to_color(c) for c in Color(\"red\").range_to(\"blue\", size)])\n\n    # Draw nodes and calculate node forces\n    @jax.vmap\n    def dots(p, i, color):\n        d = circle(0.1).translate(p[0], p[1]).fill_color(color)\n        return d, np.abs(x).sum()\n\n    out, gravity = dots(x, np.arange(size), color)\n\n    # Draw nodes and calculate edge forces\n    @jax.vmap\n    def connect(inp, out):\n        a, b = x[inp], x[out]\n        # `make_path`` is a helper around `seg`.\n        return make_path([(a[0], a[1]), (b[0] + eps, b[1] + eps)]), L2(a - b)\n\n    lines, tension = connect(nodes, edges)\n\n    # Style and compose graph.\n    out = out.with_envelope(empty()).line_width(2)\n    # `with_envelope` speeds up rendering, `rectangle` is a helper for paths\n    out = (\n        rectangle(5, 5).fill_color(\"white\")\n        + lines.with_envelope(empty()).line_width(1)\n        + out\n    )\n\n    # Apply forces\n    repulse = ((1 / (1e-3 + L2(x[:, None] - x))) * (1 - np.eye(size))).sum()\n\n    score = a * gravity.sum() + b * repulse + c * L2(tension - spring)\n\n    return score, out.layout(500)\n```\n\nGraphs where there is a single hub and several surrounding nodes.\n\n```\nsize = 50\naround = 5\ngroups = size // around\nmatrix = jax.random.uniform(jax.random.PRNGKey(0), (size, 2)) * 2 - 1\nedges = (np.arange(size) // around) * around\nopt(matrix, partial(force_directed, edges=edges), steps=500)\n```\n\nGraphs where each is a ring.\n\n```\nsize = 100\nmatrix = jax.random.uniform(jax.random.PRNGKey(0), (size, 2)) * 2 - 1\nedges = (np.arange(size) + 1) % 10 + (np.arange(size) // 10) * 10\nopt(matrix, partial(force_directed, edges=edges), steps=500)\n```\n\nUp until this point we have been entirely focused on the compute the vector diagram from its parameters.\n\n\\[\\begin{align*} \\text{vector} &= \\text{program}(x) \\\\ \\end{align*}\\]\n\nThis section considers the problem of extracting information from the vector diagram.\n\n\\[\\begin{align*} \\text{output} &= g(\\text{vector}) \\\\ \\end{align*}\\]\n\nVector diagrams are complex, but we will only need once piece of information, which is known as a *trace*. We will use the trace to to query information about the the vector diagram..\n\nA trace is a function takes in a point and a direction (a *ray*) and returns a set of distances at which the ray hits the vector diagram. This is more clear in an animation.\n\n``` python\ndef ray(t, shape, pt, v):\n    ray = seg(v).stroke().scale(t + 0.1).translate_by(pt) + circle_at(pt, 0.1)\n    trace = shape.get_trace()\n    # Length to hit shape.\n    distances, mask = trace(pt, v)\n    return shape + ray + circle_at(np.max(distances) * v + pt, 0.1).fill_color(\"red\")\n\nshape = circle_at(P2(0, 0), 1).fill_color(\"white\")\nanimate(partial(ray, shape=shape, pt=P2(0.2, -0.1), v=V2(2, 1)))\n```\n\nRecall that arcs are circle transformed by some \\(A\\) (with an angle range). This implies that the trace function can applying the inverse of this transform \\(A^{-1}\\) to the ray, and calcuting its distance to the unit circle.\n\nTo calculate this value we need to do some high-school algebra. The formula for this calculation is to find a common point on the ray and the unit circle. Let \\(r\\) be the length of the trace.\n\n\\[ \\begin{align*} x = v_x r + p_x & y = v_y r + p_y \\\\ x^2 + y^2 &= 1 \\\\ \\Rightarrow& (v_x^2 + v_y^2) r^2 + 2(v_x p_x + v_y p_y) r + (p_x^2 + p_y^2 - 1) = 0 \\\\ \\Rightarrow& \\| v \\| r^2 + 2 p \\cdot v r + (\\| p \\| - 1) = 0 \\\\ \\end{align*} \\]\n\nWe then apply the quadratic formula to solve for roots.\n\n\\[\\begin{align*} \\frac{-b \\pm \\sqrt{b^2 - 4a c}}{2a} \\end{align*}\\]\n\nThe two values correspond to the two possible intersections. Negative roots correspond to missing the circle entirely. Once we find the value, we can also compute the angle of the intersection and check if it is within the arc.\n\n```\nshape = arc_seg(V2(2, 0), 0.5).stroke().translate(-1, 0)\nanimate(partial(ray, shape=shape, pt=P2(0.2, 0.3), v=V2(1, -1)))\n```\n\nNote that there may be multiple intersections between the ray and the shape. This means trace needs to return a set of values. Working with sets is difficult in languages like Jax because array size needs to remain constant in repeated calls to a function. To handle this issue we return a fixed size array with a mask.\n\nHere’s an example that shows multiple intersections, as well as use of complex affine transformations.\n\n``` python\ndef draw_line(p1, p2):\n    \"Helper function to draw a line between two points.\"\n    return Path.from_points([p2, p1 + eps]).stroke()\n\ndef draw_trace(diagram, pt, v):\n    \"Draw a trace of a vector diagram.\"\n    trace = diagram.get_trace()\n    distances, mask = trace(pt, v)\n\n    # Draw the raw and trace.\n    line = seg(V2(0, -1 / 2)).stroke()\n    out = (diagram + draw_line(pt, pt+v) + circle_at(pt, 0.1).fill_color(\"white\")).line_width(2)\n    pos = (1 - mask[:, :])[..., None, None] * 100 + distances[..., None, None] * v + pt\n    out = out + draw_line(pos[:, 0], pos[:, 1]).line_color(\n        \"white\"\n    ).line_width(10)\n    for i in range(distances.shape[1]):\n        out = out + circle_at(pos[:, i], 0.1).fill_color(\"red\").line_width(1 * mask[:, i]).fill_opacity(mask[:, i])\n    return out, distances\n\ndef satelite(d, t, px=0, py=0):\n    # Place the ray\n    d = d.center_xy()\n    affine = tx.rotation(t * 2 * np.pi)\n    pt, v = P2(0, 1.5), V2(0, -1.5)\n    pt, v = affine @ pt + V2(px, py), affine @ v\n\n    return draw_trace(d, pt, v)[0]\n\ndef crescent():\n    part1 = arc_seg(unit_x, 0.5)\n    part2 = arc_seg(-unit_x, -0.2)\n    return part1, part2, (part1 + part2).close().stroke()\n\nanimate(partial(satelite, crescent()[2]), lw=False)\n```\n\nInternally, we are storing all the arcs that make up a path. each of these is reset to the unit circle where we then apply a ray-circle formula to find possible intersections.\n\nArc 1\n\n```\nanimate(partial(satelite, crescent()[0].stroke()), lw=False)\n```\n\nArc 2.\n\n```\nanimate(partial(satelite, crescent()[1].stroke(), py=-0.2), lw=False)\n```\n\nNow here’s the cool part. Because all of these traces are computed with using differentiable transformations (Section 1), and they are in Jax (Section 2), we can optimize through them. Here we set up a function that given any shape and ray will compute and draw the trace.\n\nUsing this function, we can then optimize for trace properties. Here we search for the outer angle that leads to the smallest gap between trace distances.\n\n``` python\ndef trace_width(d, t):\n    d = d.center_xy()\n    # Outer satelite\n    affine = tx.rotation(t)\n    pt, v = affine @ P2(0, 1.5), affine @ V2(0, -1.5)\n\n    # Trace the image\n    out, p = draw_trace(d, pt, v)\n    score = (p[:, 1] - p[:, 0])\n    return score.sum(), (grid + out.with_envelope(empty())).layout(500)\n\nopt(np.array(np.pi / 4.0), partial(trace_width, crescent()[2]), rate=0.1)\n```\n\n*Rasterization* is the process of converting from a vector diagram to a visible image, specifical a height \\(\\times\\) width \\(\\times\\) color array. In our formula it represents this step.\n\n\\[\\begin{align*} \\text{image} &= \\text{rasterize}(\\text{vector}) \\\\ \\end{align*}\\]\n\nHaven’t we been doing this all along though? Well kind of. We have been cheating under the hood and using [Cairo](https://www.cairographics.org/) an open-source rasterizer. We would like to do it ourself in Jax.\n\nTo do rasterization, we will use a [ Scanline](https://en.wikipedia.org/wiki/Scanline_rendering) algorithm. We are choosing this approach because it is relatively simple compared to other methods. We are motivated by the paper\n\nThe algorithm works by walking down each row one pixel at a time, runs a trace, and then fills in the middles. We will use an *even-odd* rule for rasterizing where we consider all pixels after every other intersection to be inside.\n\n``` python\ndef scanline(t, shape):\n    return draw_trace(shape, P2(1, 1 + 10 * t), V2(1, 0))[0]\n\nshape = crescent()[2].scale(7).rotate(70).translate(6, 8).fill_color(\"orange\")\nanimate(partial(scanline, shape=shape), grid=hgrid, steps=90)\n```\n\nHere’s the critical rasterization code. This function produces 1 row of the image given the trace for that row. Its job:\n\nHere’s how you would do this in simple python. We use the term split to mean the distance of the trace from the left-hand side of the picture. Mask is 1 if a split is real and 0 if it is not.\n\n``` python\ndef render_line_python(splits, mask):\n    import math\n    splits_int = {}\n    even = True\n    \n    # Discretize each split to its pixel cell.\n    for mask, split in zip(mask, splits):\n       if mask:\n           splits_int[math.floor(split)] = split\n\n    # Create a blank row. \n    scene = [0] * 100   \n\n    for j in range(100):\n       if j in splits_int:\n           split = splits_int[j]\n            # Set boundaries based on where pixel fell inside.\n           if even:\n               scene[j] = 1 - (split - j)\n           else:\n               scene[j] = split - j\n\n            # Compute the even / odd position of each split\n           even = not even\n       else:\n           if not even:\n                # Fill in inside.\n               scene[j] = 1\n    return np.array(scene)\n\nrender_line_python([3.4, 9.7, 16.9], [1, 1, 0])\nArray([0. , 0. , 0. , 0.6, 1. , 1. , 1. , 1. , 1. , 0.7, 0. , 0. , 0. ,\n       0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ,\n       0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ,\n       0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ,\n       0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ,\n       0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ,\n       0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ,\n       0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ], dtype=float64)\n```\n\nIt’s a relatively simple algorithm, but remember we can’t use conditionals, mutation, or dictionaries in Jax. And we want it to run in parallel!\n\nWe will instead modify this approach to do scanline rendering in parallel. The trick will be to first mark where we cross a boundary and then use a *cumulative sum* to combine them together. Here’s the same algorithm in Jax.\n\n``` python\nSIZE = 100\nblank = np.ones((SIZE, SIZE, 3))\n\n@jax.jit\n@jax.vmap\ndef render_line(splits, mask):\n    # Discretize each split to its pixel cell.\n    split_int = np.floor(splits).astype(int)\n\n    # `np.where`` will be our if statement replacement.\n    # Set `ind` to these values, unless it is masked.\n    ind = np.where(mask, split_int, 1000)\n\n    # Create a blank row.\n    zero = np.zeros(SIZE)\n\n    # Compute the even / odd position of each split\n    # (Assume sorted with masked last)\n    loc = np.arange(splits.shape[-1]) % 2\n\n    # Set even->odd pixels as 1 and odd->even as 0.\n    inout = np.where(loc, -1, 1)\n    scene = zero.at[ind].add(inout)\n\n    # Run a cumulative sum to fill in inside.\n    scene = np.cumsum(scene, axis=-1)\n\n    # Set boundaries based on where pixel fell inside.\n    scene = scene.at[ind].set(np.where(loc, 0, 1) - inout * (splits - split_int))\n    \n    # If we didn't end on even, something went wrong.\n    return np.where(mask.sum() % 2 == 0, scene, zero)\n\nrender_line(np.array([[3.4, 9.7, 16.9]]), np.array([[1, 1, 0]]))\nArray([[0. , 0. , 0. , 0.6, 1. , 1. , 1. , 1. , 1. , 0.7, 0. , 0. , 0. ,\n        0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ,\n        0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ,\n        0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ,\n        0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ,\n        0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ,\n        0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ,\n        0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ]], dtype=float64)\nshape = crescent()[2].scale(70).rotate(70).translate(60, 80).fill_color(\"orange\")\nplt.imshow(render_line(*shape.get_trace()(P2(0, np.arange(100)), V2(1, 0))))\nNone\n```\n\nNice! 10 lines of code and we have a rasterizer. It works pretty nice off the bat. The main problem though is that it doesn’t yet look as good as Cairo’s version. The textbook issue is that we have *aliasing*. If you zoom into the edge of our moon you will see jagged little edges.\n\nWe’ll apply convolution smoothing on each row independently. The idea is to set the value of each pixel based on the surrounding values from the vector diagram. This is a 1D simplification but it makes the math much easier. The formula is,\n\n\\[I(x) = \\int k(u) f(x- u; \\theta) du\\]\n\nWhere \\(f(x-u)\\) is the value of the line our vector image, \\(k(x)\\) is a *kernel*, and \\(I(x)\\) is our output image. The *kernel* is just a function that weights the contribution of neighbors based on how far away they are.\n\nWe can approximate this integral with Monte-Carlo sampling.\n\n\\[ I(x) \\approx \\frac{1}{N} \\sum_i k(u_i) f(x - u_i) \\]\n\nWhere \\(u_i\\) are \\(N\\) chosen points around \\(x\\).\n\nWe have two choices here: the kernel function and the monte-carlo scheme.\n\nFrom an ML perspective, you can think of this a just applying a 1D convolution to each row of the image with a predefined kernel.\n\n``` python\n# Kernel width\nkern = 11\nsamples = np.arange(kern) - (kern // 2)\n\ndef kernel(offset):\n    off_samples = samples - offset\n    kernel = kern - np.abs(off_samples)\n    return np.maximum(0, kernel / (kern - np.abs(samples)).sum())\n\n# Allow us to offset the grid. \nplt.plot(samples, kernel(0), \"o-\")\nplt.plot(samples - 0.2, kernel(0.2), \"o-\")\nplt.plot(samples + 0.5, kernel(0.5), \"o-\")\n[<matplotlib.lines.Line2D at 0x74204423f4a0>]\n```\n\nCompare the new image to the original. Note the little “hairs” that come off on both sides. These will become important in the next section.\n\n``` python\n@jax.vmap\ndef convolve(line):\n    k = kernel(0)\n    return jax.vmap(lambda s: line[s + samples] @ k)(np.arange(line.shape[0]))\n\nplt.imshow(convolve(render_line(*shape.get_trace()(P2(0, np.arange(100)), V2(1, 0)))))\nNone\n```\n\nNow let’s put everything together. There are three steps.\n\nWe will do it twice once for rows and once for columns.\n\n```\ndirections = list(map(np.stack, [(V2(0, 1), V2(1, 0)), (V2(1, 0), V2(0, 1))]))\n\n@jax.jit\ndef render_shape(img, s):\n    S = np.arange(SIZE)\n    trace = s.get_trace()\n\n    def render(img, d):\n        pt, v = d\n        # 1. Scan\n        raster = render_line(*trace(tx.to_point(pt * S[:, None, None]), v))\n        # 2. Apply a 1d conv\n        raster = convolve(raster)\n        # 3. Compose\n        img = (1 - raster[..., None]) * img + raster[..., None] * s.style.fill_color_\n        return np.swapaxes(img, 0, 1), raster\n\n    return jax.lax.scan(render, img, directions)\n\nplt.imshow(render_shape(blank, shape)[0])\nNone\n```\n\nWe apply the same process for all the parts of our diagram.\n\n``` python\n@jax.vmap\ndef make_shape(i):\n    r, g = to_color(\"red\"), to_color(\"green\")\n    return shape.translate(-30 + i * 40, -10 + i * 30).fill_color(r * i + g * (1 - i))\n\ndef render_shapes(shapes):\n    return jax.lax.scan(render_shape, blank, shapes)[0]\n\nplt.imshow(render_shapes(make_shape(np.arange(5) / 5)))\nNone\n```\n\nAnd then just a check that we can render for all transforms.\n\n``` python\n@jax.jit\ndef im(t):\n    s = shape.center_xy().rotate_by(t).translate(50, 50)\n    return render_shape(blank, s)[0]\n\nos.system(\"rm /tmp/test*.png 2> /dev/null\")\nfor t in range(T):\n    plt.imshow(im(t / T))\n    plt.savefig(f\"/tmp/test.{t:02d}.png\")\n    plt.clf()\n\nanimate_out(sorted(glob.glob(\"/tmp/test.*.png\")))\n<Figure size 640x480 with 0 Axes>\n```\n\nThis rasterizer is fully end-to-end. This now allows use to Now we can do our first optimization.We need to specify a goal target image, and then render our own guess. Here we start with a blue image and target an orange one. We then optimize through the process (using derivatives) to find the color.\n\n``` python\ngoal = render_shape(blank, shape)[0]\n\ndef make_shape(color):\n    return shape.fill_color(color)\n\ndef loss(x):\n    y = render_shape(blank, x)[0]\n    return np.pow(y - goal, 2).sum(), ((hgrid + x) | (hgrid + shape)).layout(500)\n\nopt(to_color(\"blue\"), lambda x: loss(make_shape(x)), rate=0.1)\n```\n\nIn this section we complete the full pipeline, and differentiate through both steps.\n\n\\[\\begin{align*} \\text{vector} &= \\text{program}(x) \\\\ \\text{image} &= \\text{rasterizer}(\\text{vector}) \\\\ {\\cal L}(\\text{image}, \\text{target}) &= \\| \\text{image} - \\text{target} \\| \\\\ \\end{align*}\\]\n\nIn fact, the last example shows we can do this already. We make a vector diagram, rasterize it and then just optimize.\n\nUnfortunately though if the images didn’t line up exactly things go wrong. The problem is that a key step in our rasterization has a derivative of zero. Take a look at the `render_line`\n\ncode and see if you can spot it.\n\n```\nnew_shape = shape.translate(15, 35)\ngoal = render_shape(blank, new_shape)[0]\n\ndef make_shape(color):\n    return shape.fill_color(color)\n\ndef loss(x):\n    y = render_shape(blank, x)[0]\n    return np.pow(y - goal, 2).sum(), ((hgrid + x) | (hgrid + new_shape)).layout(500)\n\nopt(to_color(\"orange\"), lambda x: loss(make_shape(x)), rate=0.1)\n```\n\nThe problem is that we are never sending a derivative back through the `split`\n\nvariable to the traces, i.e. \\(\\frac{d {\\cal L}}{d \\ \\text{trace}}=0\\) That is because the only way split is used is as a discrete index into the array.\n\nTo do step, we are going to have to find a workaround.\n\n[Differentiable Vector Graphics Rasterization for Editing and Learning](https://cseweb.ucsd.edu/~tzli/diffvg/) describes an elegant method for producing derivatives. We’ll use a simplified version here which demonstrates the core idea.\n\nThe trick is to take advantage of the anti-aliasing step. Recall the formula,\n\n\\[I(x) = \\int k(u) f(x- u) du\\]\n\nWe would like to take the derivative of this function with respect to the trace.\n\n\\[\\frac{d}{d \\ \\text{trace}} \\int k(u) f(x- u) du\\]\n\nBut where is the trace in this function? It is used to produce the splits, i.e. the *inside* and *outside* parts where \\(f\\) changes. We call each of these independent regions \\(A_i\\).\n\n\\[\\sum_i \\frac{d}{d \\ \\text{trace}} \\int_{A_i} f(x-u) k(u) du\\]\n\nThese intersection points does impact where the *boundary* of the integral is. To optimize positioning we need a signal from this boundary position.\n\nThe methodology for differentiating through integral boundaries is an identity called [Leibniz’s integral rule](https://en.wikipedia.org/wiki/Leibniz_integral_rule).\n\n\\[\\frac{d}{db} \\int_{0}^{b} z(t, b) dt = z(b, b) + \\int_{0}^{b} \\frac{d z(t, b)}{d b} dt\\]\n\nThe second part is just within the boundaries. We have taken care of the already just naturally through our convolution above. This is why the color-match optimization works already.\n\nThe first part though we need to handle ourselves. For our simplified problem this takes a nice form since \\(z(b, b) = f(x - \\text{split})k(\\text{split})\\).\n\n\\[f(x-\\text{split}) k(\\text{split}) \\approx (f(x-\\text{split} + \\epsilon) - f(x-\\text{split} - \\epsilon))k(\\text{split})\\]\n\nWhere the approximation here is necessary because by definition \\(f(x - \\text{split})\\) falls directly on a boundary point. Intuitively the math is just telling us to move boundaries closer/farther to targets proportional to their kernel value.\n\nSince to implement backprop we need, \\[\\frac{d L}{d I} \\frac{d I}{d \\text{trace}}\\]\n\nThe term \\(\\frac{d L}{d I}\\) is the size of the image, so the implementation looks like applying the reverse of the triangle kernel from above to each around of the split points to this array.\n\n``` python\n@jax.custom_vjp\ndef boundary(scene, splits, mask):\n    \"Do nothing on the forward pass.\"\n    return scene\n\ndef f_fwd(scene, splits, mask):\n    \"Blank forward pass.\"\n    return scene, (scene, splits, mask)\n\ndef f_bwd(res, g):\n    \"Backwards pass applies Leibniz.\"\n    f, splits, mask = res\n    split_int = np.floor(splits).astype(int)\n\n    def grad_p(s, s_off):\n        \"Compute the kernel values of neighbors around this point\"\n        off = s_off - s\n        # Chain rule part, g is dL / dI. Apply kernel around split s. \n        v = g[s + samples]\n        return (v * (f[s + 1] - f[s - 1])).sum(-1) @ kernel(off)\n\n    # For each split compute dL/dI\n    r = jax.vmap(grad_p, in_axes=(-1, -1))(split_int, splits) * mask\n    r = np.where(mask.sum() % 2 == 0, r, np.zeros(splits.shape))\n    \n    return g, -r, None\n\nboundary.defvjp(f_fwd, f_bwd)\n```\n\nWe then can apply it on top of our rendering function. The idea is that this should not change the image in any way, but it allows us to correct the derivative.\n\n``` python\n@jax.jit\ndef boundary_shape(img, s):\n    S = np.arange(SIZE)\n    trace = s.get_trace()\n\n    def bound(img, d):\n        pt, v = d\n        s, m = jax.vmap(lambda i: trace(tx.to_point(pt * i), v))(S)\n        bound = lambda s, m: jax.vmap(lambda s, m, im: boundary(im, s, m))(s, m, img)\n        img = bound(s, m)\n        return np.swapaxes(img, 0, 1), None\n\n    return jax.lax.scan(bound, img, directions)[0], None\n\ndef boundaries(img, shapes):\n    return jax.lax.scan(boundary_shape, img, shapes)[0]\n```\n\nAnd that’s it! We have a differentiable rasterizer. A bit of work, but only really one custom derivative and the rest we got from the framework for free.\n\nThis first one shows we can deal with occulusion and multiple shapes for free.\n\n``` python\n@jax.vmap\ndef make_shape1(p):\n    loc, color = p\n    return (\n        shape.center_xy().rotate_by(1 / 4).translate(loc[0], loc[1]).fill_color(color)\n    )\n\n# Target vector / raster images\ngoal_shapes = make_shape1(\n    (np.array([[50, 40], [40, 50]]), np.stack([to_color(\"blue\"), to_color(\"orange\")]))\n)\n\n# Target vector / raster images\ngoal_shapes = make_shape1(\n    (np.array([[50, 40], [40, 50]]), np.stack([to_color(\"blue\"), to_color(\"orange\")]))\n)\ngoal = render_shapes(goal_shapes)\n\n# Initial parameters $x$\nstart_param = (\n    np.array([[60, 30], [60.0, 60.0]]),\n    np.stack([to_color(\"green\"), to_color(\"green\")]),\n)\n\n# L2 Loss and draw\ndef loss(goal, goal_shapes, x):\n    y = render_shapes(x)\n    y = boundaries(y, x)\n    return np.pow(y - goal, 2).sum(), (\n        bgrid\n        + concat(x).with_envelope(empty())\n        + concat(goal_shapes).with_envelope(empty()).fill_opacity(0.2)\n    ).layout(500)\n\nopt(\n    start_param,\n    lambda x: loss(goal, goal_shapes, make_shape1(x)),\n    steps=500,\n    save_every=5,\n    rate=0.1,\n)\n```\n\nThis one parameterizes the structure of the arcs that make up the path.\n\n``` python\n@jax.vmap\ndef make_shape2(params):\n    s, l, r = params\n    s = np.where(np.abs(s) < 1e-3, 1e-3, s)\n    part1 = arc_seg(unit_x, s[0])\n    part2 = arc_seg(unit_y, s[1])\n    part3 = arc_seg(-unit_x, s[2])\n    part4 = arc_seg(-unit_y, s[3])\n    d = (part1 + part2 + part3 + part4).close().stroke()\n    shape = d.scale(12 + r[0]).rotate(70)\n    return shape.center_xy().rotate_by(0).translate(50, 30).fill_color(\"blue\")\n\ngoal_params = (\n    np.array([[0.5, 0.2, 0.7, -0.9]]),\n    np.array([[50.0, 30]]),\n    np.array([[5.0]]),\n)\ngoal_shapes = make_shape2(goal_params)\ngoal = render_shapes(goal_shapes)\nstart_param = (np.array([[-0.5, 0.1, 2, 1]]), np.array([[50.0, 30]]), np.array([[0.0]]))\n\nopt(\n    start_param,\n    lambda x: loss(goal, goal_shapes, make_shape2(x)),\n    steps=1000,\n    save_every=5,\n    rate=0.01,\n)\n```\n\nIn this example we create a bunch or red and blue circles with vmap and then try to fill them in.\n\n``` python\n@jax.vmap\ndef make_shape3(params, i):\n    s, l = params\n    s = np.minimum(np.maximum(s, 5.0), 15)\n    red, blue = to_color(\"red\"), to_color(\"blue\")\n    return (\n        circle_at(tx.X.origin, s)\n        .translate(l[0], l[1])\n        .fill_color(np.where(i % 2, red, blue))\n    )\n\nmake_shape3 = partial(make_shape3, i=np.arange(10.0))\n\ngoal_param = (\n    np.array([10 + random.random() for i in range(10)]),\n    np.array([20 + random.random() * 50 for i in range(20)]).reshape(10, 2),\n)\ngoal_shapes = make_shape3(goal_param)\n\nmake_shape3 = partial(make_shape3, i=np.arange(20))\n\ngoal = render_shapes(goal_shapes)\nstart_param = (\n    np.array([10 + random.random() for i in range(20)]),\n    np.array([20 + 50 * random.random() for i in range(40)], np.float64).reshape(20, 2),\n)\n\nopt(\n    start_param,\n    lambda x: loss(goal, goal_shapes, make_shape3(x)),\n    steps=1000,\n    save_every=5,\n    rate=0.1,\n)\n```\n\nIn this exmaple we render a star with a gap in the middle to test the inside/outside aspect of the renderer.\n\n``` python\n @jax.vmap\n def star(x):\n     sides = 5\n     edge = Trail.hrule(1)\n     return Trail.concat(\n        edge.rotate_by((2 * i) / sides) for i in range(sides)\n     ).close().stroke().center_xy().scale(x[0]).rotate(x[1]).translate(x[2], x[3] ).fill_color(\"orange\")\nrandom.seed(1)\n\nstart_param = np.array([[65.,random.random(),  40 + i * 20, 40 + 20 * j] for i in range(1) for j in range(1)])\ng = np.array([[85., 50, 50 + i * 20, 50 + 20 * j] for i in range(1) for j in range(1)])\ngoal_shapes = star(g)\ngoal = render_shapes(star(g))\n \nopt(\n    start_param,\n    lambda x: loss(goal, goal_shapes, star(x)),\n    steps=500,\n    save_every=2,\n    rate=0.35,\n    verbose=False,\n)\n```\n\nIn this exmaple we render a set of gears that need to fit together.\n\n``` python\n@jax.vmap\ndef gear(x):\n    sides = 5\n    edge = Trail.hrule(1)\n    return (\n        Trail.concat(edge.rotate_by((i) / sides) for i in range(sides))\n        .close()\n        .stroke()\n        .center_xy()\n        .scale(15)\n        .rotate_by(x[1])\n        .translate(25 + x[2], 25 + x[3])\n        .fill_color(\"orange\")\n    )\n\nstart_param = np.array(\n    [\n        [15.0, random.random(), 5 + i * 20, 10 + 20 * j]\n        for i in range(3)\n        for j in range(3)\n    ]\n)\ng = np.array(\n    [\n        [15.0, random.random(), 1 + i * 20, 1 + 20 * j]\n        for i in range(3)\n        for j in range(3)\n    ]\n)\ngoal_shapes = gear(g)\ngoal = render_shapes(gear(g))\n\nopt(\n    start_param,\n    lambda x: loss(goal, goal_shapes, gear(x)),\n    steps=500,\n    save_every=2,\n    rate=0.05\n)\n```\n\nAll this work brings us to our final goal, rendering an actual image. Here we take an image from the web and convert it to an array. Then we start with a bunch of ellipses (scale-transformed circles) and see how well they can do to approximate it. Results are pretty neat!\n\nGrab an image from the web and covert to RGB.\n\n```\ngoal = pix\n\n# Draw a bunch of circles.\n# Parameterize by scale, translation, and roation.\n@jax.vmap\ndef make_smiley(p):\n    s, l, r, c, rot = p\n    s = np.where(np.abs(s) < 1e-3, (np.abs(s) / s) * 1e-3, s)\n    d = circle_at(P2(0, 0), 1).center_xy()\n    r = np.where(np.abs(r) < 1e-3, 0.1, (np.abs(r) / r) * r)\n    return (\n        d.rotate(rot)\n        .scale(1)\n        .scale_x((r[0]))\n        .scale_y((r[1]))\n        .translate(l[0], l[1])\n        .fill_color(jax.nn.sigmoid(c))\n    )\n\nrandom.seed(1)\n\n# Draw 20 circles with random starting\n# values.\nn = 20\nstart_param = (\n    np.zeros((n, 4)) + 0.1,\n    np.array(\n        [[20 + 60 * random.random(), 20 + 60 * random.random()] for i in range(n)]\n    ),\n    np.ones((n, 2)) + 2 * np.arange(n + 1, 1, -1)[:, None],\n    np.ones((n, 3)),\n    np.zeros((n, 1)),\n)\n\n# Reduce color precision of the goal slightly.\ndef reduce(y):\n    return np.floor(y * 40) / 40\n\ndef loss(goal, x):\n    y = render_shapes(x)\n    y = boundaries(y, x)\n    return np.pow(y - reduce(goal), 2).sum(), None\n\n@jax.jit\ndef show(x):\n    return (bgrid + concat(make_smiley(x)).with_envelope(empty()).line_width(0)).layout(\n        500\n    )\n\n# Optimize!\nopt(\n    start_param,\n    lambda x: loss(goal, make_smiley(x)),\n    show=show,\n    steps=1000,\n    save_every=10,\n    rate=0.1,\n)\n```\n\nSo the goal of this post was to understand what differentiable programming really look like for challening problems. Of course graphics is a particularly nice area for this kind of silliness. Since everything has a nice and well-defined mathematical form, with enough work we can push through the necessary gradients. We only really had to use one major trick to make derivative work, other wise it was just converting if’s and for’s to np.where and vmap.\n\nStill I think this sort of method is still really underexplored in other areas. We have really nice tools for writing very complex programs in a fully vectorized way, and yet we aren’t really doing it yet as much as feels possible. I hope I can convince you that while a little scary, this kind of programming is pretty powerful and realistically accessible for a lot of different domains.\n\n-Sasha", "url": "https://wpnews.pro/news/a-tour-of-differentiable-rasterization", "canonical_source": "https://srush.github.io/DiffRast/", "published_at": "2026-08-26 07:06:53+00:00", "updated_at": "2026-08-26 07:44:18.655549+00:00", "lang": "en", "topics": ["machine-learning", "computer-vision", "ai-research", "developer-tools"], "entities": ["Sasha Rush", "Dan Oneață", "Chalk", "Jax", "PyTorch"], "alternates": {"html": "https://wpnews.pro/news/a-tour-of-differentiable-rasterization", "markdown": "https://wpnews.pro/news/a-tour-of-differentiable-rasterization.md", "text": "https://wpnews.pro/news/a-tour-of-differentiable-rasterization.txt", "jsonld": "https://wpnews.pro/news/a-tour-of-differentiable-rasterization.jsonld"}}