I'm a bit skeptical about the implementation.

For one thing I don't think it really needs a registry of operation-to-operation mappings. At the moment it's essentially a layer of checking for "is this jaxpr correct". This can be done without operation substitution, by simply analysing the jaxpr and raising an error if e.g. two disparate units are added together.

For another I don't think this will scale well: if I call @units on a function that itself calls @units then the inner one will have units_jaxpr evaluated on it multiple times, which isn't necessary and will probably explode if one has a heavily @units-decorated codebase. Probably nested @units calls should seamlessly join up into a single level of checking (somehow), and each time they encounter each other they should make sure that the units being propagated by the outer level agree with the units annotated on the inner level.

(+various nits: **kwargs isn't passed to units_jaxpr; what if an argument has a magnitude attribute; for an input_units API I'd suggesting copying the API Equinox uses for e.g. filter_vmap which uses inspect.signature and signature.bind_partial etc.)

For one thing I don't think it really needs a registry of operation-to-operation mappings. At the moment it's essentially a layer of checking for "is this jaxpr correct". This can be done without operation substitution, by simply analysing the jaxpr and raising an error if e.g. two disparate units are added together.

Not really! It's actually also doing unit conversions in the cases when that's necessary. It would be possible to provide a more minimal API that just checks, but that's not what I need!

For another I don't think this will scale well: if I call @units on a function that itself calls @units then the inner one will have units_jaxpr evaluated on it multiple times, which isn't necessary and will probably explode if one has a heavily @units-decorated codebase. Probably nested @units calls should seamlessly join up into a single level of checking (somehow), and each time they encounter each other they should make sure that the units being propagated by the outer level agree with the units annotated on the inner level.

Totally agreed - I'd love to hear suggestions for approaches to this!