CustomVjp

Define custom Vision-Jacobian Product (backward pass) for valid differentiation.

Abstract Signature:

CustomVjp(fun: Callable, nondiff_argnums: tuple = ())

PyTorch

API: torch.autograd.function.Function
Strategy: Direct Mapping
Official Docs ↗

JAX (Core)

API: jax.custom_vjp
Strategy: Direct Mapping
Official Docs ↗

TensorFlow

API: tf.custom_gradient
Strategy: Direct Mapping
Official Docs ↗

Flax NNX

API: flax.nnx.custom_vjp
Strategy: Direct Mapping
Official Docs ↗

PaxML / Praxis

API: jax.custom_vjp
Strategy: Direct Mapping
Official Docs ↗