def f_joint(*args):
    """This function takes all args, first counter+1 are non-diff ones."""
    nondiff_args = list(args[:counter+1])
    def f_diff(*diff_args):  # Takes only diff args, will define custom grad.
      args = nondiff_args + list(diff_args)
      return f(*args)
    def f_vjp(*diff_args):  # Custom VJP for diff args.
      args = nondiff_args + list(diff_args)
      out, residual = f_fwd(*args)
      def vjpfn(g):
        bwd_args = [residual, g]
        res = f_bwd(*bwd_args)
        return res[counter+1:]
      return out, vjpfn
    # This is the function taking only diff args with custom vjp.
    f_diff_vjp = backend()['custom_grad'](f_vjp, f_diff)
    # Call it on the diff args.
    return f_diff_vjp(*args[counter+1:])
This function takes all args, first counter+1 are non-diff ones.
trax/fastmath/ops.py