From 9f4fd0573dc8b53165b6dd109b4cc2b5bfcd5b7b Mon Sep 17 00:00:00 2001 From: Nathaniel Starkman Date: Sat, 25 Jul 2026 11:00:07 -0400 Subject: [PATCH] =?UTF-8?q?Backport=20PR=20#801:=20=E2=99=BB=EF=B8=8F=20re?= =?UTF-8?q?factor:=20share=20unit=20strip/re-apply=20plumbing=20across=20e?= =?UTF-8?q?xperimental=20AD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/unxt/_src/experimental.py | 169 +++++++++++++++++----------------- 1 file changed, 84 insertions(+), 85 deletions(-) diff --git a/src/unxt/_src/experimental.py b/src/unxt/_src/experimental.py index 3cf079d2..e7ff0137 100644 --- a/src/unxt/_src/experimental.py +++ b/src/unxt/_src/experimental.py @@ -102,6 +102,81 @@ def where[R: AbstractQuantity](condition: ArrayLike, x: R, y: AbstractQuantity, return replace(x, value=jax.numpy.where(condition, xv, yv)) +# =================================================================== +# Shared unit strip/re-apply plumbing for the AD wrappers below. +# The autodiff functions all strip units off the inputs, run a JAX transform on +# the magnitudes, then re-apply units. A ``None`` unit marks a plain (unitless) +# argument that is passed through untouched. + + +def _wrap_args( + args: tuple[Any, ...], theunits: tuple[AbstractUnit | None, ...], / +) -> tuple[Any, ...]: + """Re-attach units, turning each magnitude into ``Quantity(arg, unit)``.""" + return tuple( + a if un is None else Quantity(a, un) + for a, un in zip(args, theunits, strict=True) + ) + + +def _strip_args( + args: tuple[Any, ...], theunits: tuple[AbstractUnit | None, ...], / +) -> tuple[Any, ...]: + """Strip each argument to its magnitude in the corresponding unit.""" + return tuple( + a if un is None else ustrip(un, a) for a, un in zip(args, theunits, strict=True) + ) + + +def _check_int_argnums(argnums: int) -> int: + """Reject non-int ``argnums`` (only single-argument selection is supported).""" + return eqx.error_if( + argnums, + not isinstance(argnums, int), + "only int argnums are currently supported", + ) + + +def _derivative_unit(value: Any, du: AbstractUnit | None, power: int, /) -> Any: + """Output unit divided by the differentiation unit**power (``None`` -> as-is). + + ``du`` is the unit of the argument being differentiated; ``power`` is the + order of differentiation (1 for grad/jacfwd, 2 for hessian). + """ + # TODO: get Quantity[unit] / unit2 -> Quantity[unit/unit2] working + return unit_of(value) if du is None else unit_of(value) / du**power + + +def _unit_aware_jacobian[*Args, R: AbstractQuantity]( + transform: Callable[..., Any], + fun: Callable[[*Args], R], + argnums: int, + theunits: tuple[AbstractUnit | None, ...], + *, + power: int, +) -> Callable[[*Args], R]: + """Build a unit-aware ``jacfwd``/``hessian`` wrapper around ``transform``. + + ``transform`` is a JAX transform (``jax.jacfwd`` or ``jax.hessian``) whose + output keeps the function's output units, so the result unit is read back + from the transformed value and corrected by ``du**power``. + """ + argnums = _check_int_argnums(argnums) + + @ft.partial(transform, argnums=argnums) + def jacfun_mag(*args: Any) -> R: + return fun(*_wrap_args(args, theunits)) # type: ignore[arg-type] + + def jacfun(*args: *Args) -> R: + # Strip to magnitudes; they are re-wrapped into Quantities inside + # ``jacfun_mag`` before ``fun`` sees them. + value = jacfun_mag(*_strip_args(args, theunits)) + new_unit = _derivative_unit(value, theunits[argnums], power) + return type_unparametrized(value)(ustrip(value), new_unit) + + return jacfun + + def grad[*Args, R: AbstractQuantity]( fun: Callable[[*Args], R], argnums: int = 0, @@ -150,40 +225,24 @@ def grad[*Args, R: AbstractQuantity]( Quantity(Array(12., dtype=float32...), unit='m') """ + argnums = _check_int_argnums(argnums) theunits: tuple[AbstractUnit | None, ...] = tuple(map(unit_or_none, units)) # Gradient of function, stripping and adding units @ft.partial(jax.grad, argnums=argnums) def gradfun_mag(*args: Any) -> ArrayLike: - args_ = ( - (a if unit is None else Quantity(a, unit)) - for a, unit in zip(args, theunits, strict=True) - ) - return ustrip(fun(*args_)) # type: ignore[arg-type] + return ustrip(fun(*_wrap_args(args, theunits))) # type: ignore[arg-type] def gradfun(*args: *Args) -> R: - # Get the value of the args. They are turned back into Quantity - # inside the function we are taking the grad of. - args_ = tuple( # type: ignore[var-annotated] - (a if unit is None else ustrip(unit, a)) - for a, unit in zip(args, theunits, strict=True) # type: ignore[arg-type] - ) + args_ = _strip_args(args, theunits) # Evaluate the value on the same args normalized to ``units`` that the # gradient is computed from, so its unit is consistent — an input given # in a convertible unit (e.g. cm for ``units=("m",)``) yields the same - # result as the normalized unit. - qargs = tuple( # type: ignore[var-annotated] - (a if unit is None else Quantity(ustrip(unit, a), unit)) - for a, unit in zip(args, theunits, strict=True) # type: ignore[arg-type] - ) - value = fun(*qargs) + # result as the normalized unit. ``grad`` needs its own value pass + # because ``gradfun_mag`` strips the output for ``jax.grad``. + value = fun(*_wrap_args(args_, theunits)) # type: ignore[arg-type] grad_value = gradfun_mag(*args_) - # Adjust the Quantity by the units of the derivative. A dimensionless - # differentiation argument (unit ``None``) contributes no unit. - # TODO: get Quantity[unit] / unit2 -> - # Quantity[unit/unit2] working - du = theunits[argnums] - new_unit = unit_of(value) if du is None else unit_of(value) / du + new_unit = _derivative_unit(value, theunits[argnums], 1) return type_unparametrized(value)(grad_value, new_unit) return gradfun @@ -221,41 +280,8 @@ def jacfwd[*Args, R: AbstractQuantity]( Quantity(Array(12., dtype=float32...), unit='m2') """ - argnums = eqx.error_if( - argnums, - not isinstance(argnums, int), - "only int argnums are currently supported", - ) - theunits: tuple[AbstractUnit | None, ...] = tuple(map(unit_or_none, units)) - - @ft.partial(jax.jacfwd, argnums=argnums) - def jacfun_mag(*args: Any) -> R: - args_ = tuple( - (a if unit is None else Quantity(a, unit)) - for a, unit in zip(args, theunits, strict=True) - ) - return fun(*args_) # type: ignore[arg-type] - - def jacfun(*args: *Args) -> R: - # Get the value of the args. They are turned back into Quantity - # inside the function we are taking the Jacobian of. - args_ = tuple( # type: ignore[var-annotated] - (a if unit is None else ustrip(unit, a)) - for a, unit in zip(args, theunits, strict=True) # type: ignore[arg-type] - ) - # Call the Jacobian, returning a Quantity - value = jacfun_mag(*args_) - # Adjust the Quantity by the units of the derivative. A dimensionless - # differentiation argument (unit ``None``) contributes no unit. - # TODO: check the unit correction - # TODO: get Quantity[unit] / unit2 -> - # Quantity[unit/unit2] working - du = theunits[argnums] - new_unit = unit_of(value) if du is None else unit_of(value) / du - return type_unparametrized(value)(ustrip(value), new_unit) - - return jacfun + return _unit_aware_jacobian(jax.jacfwd, fun, argnums, theunits, power=1) def hessian[*Args, R: AbstractQuantity]( @@ -301,31 +327,4 @@ def hessian[*Args, R: AbstractQuantity]( """ theunits: tuple[AbstractUnit | None, ...] = tuple(map(unit_or_none, units)) - - @ft.partial(jax.hessian, argnums=argnums) - def hessfun_mag(*args: Any) -> R: - args_ = tuple( - (a if unit is None else Quantity(a, unit)) - for a, unit in zip(args, theunits, strict=True) - ) - return fun(*args_) # type: ignore[arg-type] - - def hessfun(*args: *Args) -> R: - # Get the value of the args. They are turned back into Quantity - # inside the function we are taking the hessian of. - args_ = tuple( # type: ignore[var-annotated] - (a if unit is None else ustrip(unit, a)) - for a, unit in zip(args, theunits, strict=True) # type: ignore[arg-type] - ) - # Call the hessian, returning a Quantity - value = hessfun_mag(*args_) - # Adjust the Quantity by the units of the derivative. A dimensionless - # differentiation argument (unit ``None``) contributes no unit. - # TODO: check the unit correction - # TODO: get Quantity[unit] / unit2 -> - # Quantity[unit/unit2] working - du = theunits[argnums] - new_unit = unit_of(value) if du is None else unit_of(value) / du**2 - return type_unparametrized(value)(ustrip(value), new_unit) - - return hessfun + return _unit_aware_jacobian(jax.hessian, fun, argnums, theunits, power=2)