JAX Introduces Custom Types: Enabling Differentiable Rasterization Beyond Tensors

srush_nlp · x · 2026-08-01

JAX has added support for opaque types with custom tangents. This feature allows developers to define new data types with their own identity and internal invariants, rather than just relying on transparent array containers (pytrees).

One cool use case is the proper handling of batched geometric objects. The author shared a short demo of differentiable rasterization, fitting a photo with 1500 triangles by minimizing L2 loss. This marks an interesting move for JAX beyond standard tensor operations into complex geometric structures.

Related event: JAX Introduces Custom Opaque Types for Differentiable Rendering(2 posts)→

Original post →

More from coding & agent

coding & agent channel →