Thanks for the suggestion. There's also a proposal in #10614 to make argnums more expressive in order to handle this kind of case explicitly, rather than implicitly as in your proposal. JAX's design tends to avoid implicit choices, so I suspect that's the route the team will end up taking to address this issue.

As far as I understand, #10614 will still require to split the tree into static and dynamic part.

So this won't solve the root issue of being able to differentiate only part of the tree.

Unless it was possible to mask the tree with something like:

mask = {'a': True, 'b': False}  # Differentiate with `a` but not `b`

@jax.grad(static={'x': mask})
def loss_fn(x):
  ...


loss_fn({'a': a, 'b': b})  # Differentiate with `a` but not `b`

Related to #5487.

I'm missing what's wrong with the allow_int option?

In [7]: @partial(jax.grad, allow_int=True)
   ...: def fn(vals):
   ...:   return vals['x'] + vals['y']
   ...:
   ...:
   ...: vals = {
   ...:     'x': jnp.ones((), dtype=jnp.float32),
   ...:     'y': jnp.ones((), dtype=jnp.int32),  # Fail: << Integer not supported
   ...: }
   ...: fn(vals)
Out[7]: {'x': Array(1., dtype=float32), 'y': array((b'',), dtype=[('float0', 'V')])}

I'm all for something explicit, but it should be possible to have both an explicit syntax without scarifying the user experience.

For example with something like:

@jax.grad
def loss_fn(x):
  ...


loss_fn({'a': a, 'b': jax.no_grad(b)})

It's better than having to split into 2 args because jax.no_grad() could be applied directly by the libraries (equinox,...).

@Conchylicultor Using allow_int=True does not require you to split your arguments.

I usually deal with PyTrees as input, so I resort to something like this :

import jax 
import jax.numpy as jnp 

class NoGrad:
    def __init__(self, wrapped):
        self.wrapped = wrapped
    def __repr__(self):
        return f'NoGrad({self.wrapped})'

def unwrap(x):
   return x.wrapped if isinstance(x, NoGrad) else x

jax.tree_util.register_pytree_node(
    nodetype=NoGrad,
    flatten_func=lambda x: ((), x.wrapped),
    unflatten_func=lambda x, _: NoGrad(x),
)

@jax.grad
def fn(vals):
  vals = jax.tree_map(unwrap, vals)
  return vals['x']  # y not used

vals = {
    'x': jnp.ones((), dtype=jnp.float32),
    'y': NoGrad(jnp.ones((), dtype=jnp.float32)),
}

fn(vals)
# {'x': Array(1., dtype=float32), 'y': NoGrad(1.0)}

Unless I'm missing something, this pattern is explicit and does not alter the function signature (static and dynamic splitting).

#14960 is related and I suspect could possibly help here too.