def _custom_lin(*args: TfVal, **_) -> Sequence[TfVal]:
    raise TypeError("can't apply forward-mode autodiff (jvp) to a custom_vjp function.")
def _custom_lin(*args: TfVal, **_) -> Sequence[TfVal]:
    raise TypeError("can't apply forward-mode autodiff (jvp) to a custom_vjp function.")