Add support for category corrections with gradients

[ligerlac]

Corrections of type 'Category' are currently not supported because discrete choices are inherently non differentiable. Consider the phimod example from the correctionlib documentation (link):

📈 phimod (v1)
Phi-dependent tracking efficiency, or something?
Node counts: Category: 1, Formula: 2
╭──────────── ▶ input ────────────╮ ╭──── ▶ input ────╮
│ phi (real)                      │ │ q (int)         │
│ No description                  │ │ Particle charge │
│ Range: [-inf, inf), overflow ok │ │ Values: -1, 1   │
╰─────────────────────────────────╯ ╰─────────────────╯
╭───────── ◀ output ──────────╮
│ weight (real)               │
│ Multiplicative event weight │
╰─────────────────────────────╯

Creating a correction with gradient like so

phimod = ...
my_corr = CorrectionWithGradient(phimod)

fails with ValueError: "Correction 'categorical' contains the unsupported operation type 'Category'". Despite phimod being non-differentiable w.r.t. to the discrete input q, it is, in principle, differentiable wr.t. phi.

This PR adds support to Category corrections so that they can be differentiated w.r.t. to a subset of their inputs:

my_corr = CorrectionWithGradient(phimod)  # this works now
value, grad = jax.value_and_grad(my_corr.evaluate, argnums=0)(42., -1)  # returns correctly
value, grad = jax.value_and_grad(my_corr.evaluate, argnums=1)(42., -1)  # throws a TypeError

Also supports string categorical choices (e.g. "up" / "down" in systematic uncertainties).

[nsmith-]

Can you rebase/merge the current main branch? I can't seem to do that myself.