Use @jtu.register_static as decorator on _Missing

[nstarman]

Applies jax.tree_util.register_static as a class decorator on _Missing instead of a separate post-definition call, co-locating the registration with the class.

# Before
class _Missing:
    ...

MISSING = _Missing()
jtu.register_static(_Missing)

# After
@jtu.register_static
class _Missing:
    ...

MISSING = _Missing()

---

💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.

[nstarman]

Copilot made this PR on nstarman#3 and I transferred it here.

[Copilot]

Pull request overview

This PR refactors how the internal _Missing sentinel type is registered with JAX pytrees by co-locating the registration directly on the class definition.

Changes:

• Apply @jtu.register_static as a class decorator to _Missing.
• Remove the separate post-definition pytree registration call for _Missing.

[patrick-kidger]

So I'm inclined to leave this as-is, e.g. because the older registration will work on older JAX versions.

[nstarman]

Or wait to merge until the minimum supported JAX version has register_static ?

[nstarman]

I can convert back to draft.