Autograd Utility: Batched

Autograd Utility: Batched#

Batched versions of autograd functions.

tad_mctc.autograd.batched.bjacrev(func, argnums=0, **kwargs)[source]#

Batched Jacobian of a function.

Parameters:
  • func (Callable[…, Tensor]) – The function whose result is differentiated.

  • argnums (int, optional) – The variable w.r.t. which will be differentiated. Defaults to 0.

Returns:

Batched Jacobian function.

Return type:

Callable[…, Tensor]