shrinath-suresh/blogs-docs-splitted huggingface dataset:
respect to which to take diagonal. Default: 1.
Note:
 "src" must be of the proper size in order to be embedded into
 "input". Specifically, it should have the same shape as
 "torch.diagonal(input, offset, dim1, dim2)"

Examples:
  >>> a = torch.zeros(3, 3)
  >>> a
  tensor([[0., 0., 0.],
          [0., 0., 0.],
          [0., 0., 0.]])

  >>> torch.diagonal_scatter(a, torch.ones(3), 0)
  tensor([[1., 0., 0.],
          [0., 1., 0.],
          [0., 0., 1.]])

  >>> torch.diagonal_scatter(a, torch.ones(2), 1)
  tensor([[0., 1., 0.],
          [0., 0., 1.],
          [0., 0., 0.]])
torch.nn.utils.prune.l1_unstructured
torch.nn.utils.prune.l1_unstructured(module, name, amount, importance_scores=None)
Prunes tensor corresponding to parameter called "name" in "module"
   by removing the specified amount of (currently unpruned) units
   with the lowest L1-norm. Modifies module in place (and also return
   the modified module) by:


adding a named buffer called "name+'_mask'" corresponding to the
      binary mask applied to the parameter "name" by the pruning
      method.


replacing the parameter "name" by its pruned version, while the
      original (unpruned) parameter is stored in a new parameter named
      "name+'_orig'".


Parameters:
      * module (nn.Module) -- module containing the tensor to
        prune
  * **name** (*str*) -- parameter name within "module" on which
    pruning will act.

  * **amount** (*int** or **float*) -- quantity of parameters to
prune. If "float", should be between 0.0 and 1.0 and represent
        the fraction of parameters to prune. If "int", it represents
        the absolute number of parameters to prune.
  * **importance_scores** (*torch.Tensor*) -- tensor of importance
    scores (of same shape as module parameter) used to compute
    mask for pruning. The values in this tensor indicate the
    importance of the corresponding elements in the parameter
    being pruned. If unspecified or None, the module parameter
    will be used in its place.

Returns:
      modified (i.e. pruned) version of the input module
Return type:
      module (nn.Module)
-[ Examples ]-



m = prune.l1_unstructured(nn.Linear(2, 3), 'weight', amount=0.2)
m.state_dict().keys()
   odict_keys(['bias', 'weight_orig', 'weight_mask'])
torch.lgamma
torch.lgamma(input, *, out=None) -> Tensor
Computes the natural logarithm of the absolute value of the gamma
   function on "input".
  \text{out}_{i} = \ln \Gamma(|\text{input}_{i}|)

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> a = torch.arange(0.5, 2, 0.5)
  >>> torch.lgamma(a)
  tensor([ 0.5724,  0.0000, -0.1208])
torch.gather
torch.gather(input, dim, index, *, sparse_grad=False, out=None) -> Tensor
Gathers values along an axis specified by dim.
For a 3-D tensor the output is specified by:
  out[i][j][k] = input[index[i][j][k]][j][k]  # if dim == 0
  out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1
  out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2

"input" and "index" must have the same number of dimensions. It is
   also required that "index.size(d) <= input.size(d)" for all
   dimensions "d != dim".  "out" will have the same shape as "index".
   Note that "input" and "index" do not broadcast against each other.
Parameters:
      * input (Tensor) -- the source tensor
  * **dim** (*int*) -- the axis along which to index

  * **index** (*LongTensor*) -- the indices of elements to gather

Keyword Arguments:
      * sparse_grad (bool, optional) -- If "True", gradient
        w.r.t. "input" will be a sparse tensor.
w.r.t. "input" will be a sparse tensor.
  * **out** (*Tensor**, **optional*) -- the destination tensor

Example:
  >>> t = torch.tensor([[1, 2], [3, 4]])
  >>> torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]]))
  tensor([[ 1,  1],
          [ 4,  3]])
torch.quantize_per_channel
torch.quantize_per_channel(input, scales, zero_points, axis, dtype) -> Tensor
Converts a float tensor to a per-channel quantized tensor with
   given scales and zero points.
Parameters:
      * input (Tensor) -- float tensor to quantize
  * **scales** (*Tensor*) -- float 1D tensor of scales to use,
    size should match "input.size(axis)"

  * **zero_points** (*int*) -- integer 1D tensor of offset to use,
    size should match "input.size(axis)"

  * **axis** (*int*) -- dimension on which apply per-channel
    quantization

  * **dtype** ("torch.dtype") -- the desired data type of returned
    tensor. Has to be one of the quantized dtypes: "torch.quint8",
    "torch.qint8", "torch.qint32"

Returns:
      A newly quantized tensor
Return type:
      Tensor
Example:
  >>> x = torch.tensor([[-1.0, 0.0], [1.0, 2.0]])
torch.quantize_per_channel(x, torch.tensor([0.1, 0.01]), torch.tensor([10, 0]), 0, torch.quint8)
      tensor([[-1.,  0.],
              [ 1.,  2.]], size=(2, 2), dtype=torch.quint8,
             quantization_scheme=torch.per_channel_affine,
             scale=tensor([0.1000, 0.0100], dtype=torch.float64),
             zero_point=tensor([10,  0]), axis=0)
      >>> torch.quantize_per_channel(x, torch.tensor([0.1, 0.01]), torch.tensor([10, 0]), 0, torch.quint8).int_repr()
      tensor([[  0,  10],
              [100, 200]], dtype=torch.uint8)
torch.signal.windows.blackman
torch.signal.windows.blackman(M, *, sym=True, dtype=None, layout=torch.strided, device=None, requires_grad=False)
Computes the Blackman window.
The Blackman window is defined as follows:
  w_n = 0.42 - 0.5 \cos \left( \frac{2 \pi n}{M - 1} \right) +
  0.08 \cos \left( \frac{4 \pi n}{M - 1} \right)

The window is normalized to 1 (maximum value is 1). However, the 1
   doesn't appear if "M" is even and "sym" is True.
Parameters:
      M (int) -- the length of the window. In other words, the
      number of points of the returned window.
Keyword Arguments:
      * sym (bool, optional) -- If False, returns a
        periodic window suitable for use in spectral analysis. If
        True, returns a symmetric window suitable for use in filter
        design. Default: True.
  * **dtype** ("torch.dtype", optional) -- the desired data type
of returned tensor. Default: if "None", uses a global default
        (see "torch.set_default_tensor_type()").
  * **layout** ("torch.layout", optional) -- the desired layout of
    returned Tensor. Default: "torch.strided".

  * **device** ("torch.device", optional) -- the desired device of
    returned tensor. Default: if "None", uses the current device
    for the default tensor type (see
    "torch.set_default_tensor_type()"). "device" will be the CPU
    for CPU tensor types and the current CUDA device for CUDA
    tensor types.

  * **requires_grad** (*bool**, **optional*) -- If autograd should
    record operations on the returned tensor. Default: "False".

Return type:
      Tensor
Examples:
  >>> # Generates a symmetric Blackman window.
  >>> torch.signal.windows.blackman(5)
  tensor([-1.4901e-08,  3.4000e-01,  1.0000e+00,  3.4000e-01, -1.4901e-08])

  >>> # Generates a periodic Blackman window.
Event
class torch.cuda.Event(enable_timing=False, blocking=False, interprocess=False)
Wrapper around a CUDA event.
CUDA events are synchronization markers that can be used to monitor
   the device's progress, to accurately measure timing, and to
   synchronize CUDA streams.
The underlying CUDA events are lazily initialized when the event is
   first recorded or exported to another process. After creation, only
   streams on the same device may record the event. However, streams
   on any device can wait on the event.
Parameters:
      * enable_timing (bool, optional) -- indicates if the
        event should measure time (default: "False")
  * **blocking** (*bool**, **optional*) -- if "True", "wait()"
    will be blocking (default: "False")

  * **interprocess** (*bool*) -- if "True", the event can be
    shared between processes (default: "False")

elapsed_time(end_event)
  Returns the time elapsed in milliseconds after the event was
recorded and before the end_event was recorded.
classmethod from_ipc_handle(device, handle)
  Reconstruct an event from an IPC handle on the given device.

ipc_handle()
  Returns an IPC handle of this event. If not recorded yet, the
  event will use the current device.

query()
  Checks if all work currently captured by event has completed.

  Returns:
     A boolean indicating if all work currently captured by event
     has completed.

record(stream=None)
  Records the event in a given stream.

  Uses "torch.cuda.current_stream()" if no stream is specified.
  The stream's device must match the event's device.

synchronize()
  Waits for the event to complete.

  Waits until the completion of all work currently captured in
  this event. This prevents the CPU thread from proceeding until
  the event completes.

     Note:

       This is a wrapper around "cudaEventSynchronize()": see CUDA
torch.argsort
torch.argsort(input, dim=- 1, descending=False, stable=False) -> Tensor
Returns the indices that sort a tensor along a given dimension in
   ascending order by value.
This is the second value returned by "torch.sort()".  See its
   documentation for the exact semantics of this method.
If "stable" is "True" then the sorting routine becomes stable,
   preserving the order of equivalent elements. If "False", the
   relative order of values which compare equal is not guaranteed.
   "True" is slower.
Parameters:
      * input (Tensor) -- the input tensor.
  * **dim** (*int**, **optional*) -- the dimension to sort along

  * **descending** (*bool**, **optional*) -- controls the sorting
    order (ascending or descending)

  * **stable** (*bool**, **optional*) -- controls the relative
    order of equivalent elements

Example:
  >>> a = torch.randn(4, 4)
  >>> a
  tensor([[ 0.0785,  1.5267, -0.8521,  0.4065],
[ 0.1598,  0.0788, -0.0745, -1.2700],
              [ 1.2208,  1.0722, -0.7064,  1.2564],
              [ 0.0669, -0.2318, -0.8229, -0.9280]])
  >>> torch.argsort(a, dim=1)
  tensor([[2, 0, 3, 1],
          [3, 2, 1, 0],
          [2, 1, 0, 3],
          [3, 2, 1, 0]])
CosineEmbeddingLoss
class torch.nn.CosineEmbeddingLoss(margin=0.0, size_average=None, reduce=None, reduction='mean')
Creates a criterion that measures the loss given input tensors x_1,
   x_2 and a Tensor label y with values 1 or -1. This is used for
   measuring whether two inputs are similar or dissimilar, using the
   cosine similarity, and is typically used for learning nonlinear
   embeddings or semi-supervised learning.
The loss function for each sample is:
  \text{loss}(x, y) = \begin{cases} 1 - \cos(x_1, x_2), & \text{if
  } y = 1 \\ \max(0, \cos(x_1, x_2) - \text{margin}), & \text{if }
  y = -1 \end{cases}

Parameters:
      * margin (float, optional) -- Should be a number from
        -1 to 1, 0 to 0.5 is suggested. If "margin" is missing, the
        default value is 0.
  * **size_average** (*bool**, **optional*) -- Deprecated (see
    "reduction"). By default, the losses are averaged over each
deprecated, and in the meantime, specifying either of those
        two args will override "reduction". Default: "'mean'"
Shape:
      * Input1: (N, D) or (D), where N is the batch size and D is
        the embedding dimension.
  * Input2: (N, D) or (D), same shape as Input1.

  * Target: (N) or ().

  * Output: If "reduction" is "'none'", then (N), otherwise
    scalar.
ConvReLU2d
class torch.ao.nn.intrinsic.qat.ConvReLU2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros', qconfig=None)
A ConvReLU2d module is a fused module of Conv2d and ReLU, attached
   with FakeQuantize modules for weight for quantization aware
   training.
We combined the interface of "Conv2d" and "BatchNorm2d".
Variables:
      weight_fake_quant -- fake quant module for weight
torch.vstack
torch.vstack(tensors, *, out=None) -> Tensor
Stack tensors in sequence vertically (row wise).
This is equivalent to concatenation along the first axis after all
   1-D tensors have been reshaped by "torch.atleast_2d()".
Parameters:
      tensors (sequence of Tensors) -- sequence of tensors to
      concatenate
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> a = torch.tensor([1, 2, 3])
  >>> b = torch.tensor([4, 5, 6])
  >>> torch.vstack((a,b))
  tensor([[1, 2, 3],
          [4, 5, 6]])
  >>> a = torch.tensor([[1],[2],[3]])
  >>> b = torch.tensor([[4],[5],[6]])
  >>> torch.vstack((a,b))
  tensor([[1],
          [2],
          [3],
          [4],
          [5],
          [6]])
torch.flip
torch.flip(input, dims) -> Tensor
Reverse the order of an n-D tensor along given axis in dims.
Note:
 *torch.flip* makes a copy of "input"'s data. This is different
 from NumPy's *np.flip*, which returns a view in constant time.
 Since copying a tensor's data is more work than viewing that
 data, *torch.flip* is expected to be slower than *np.flip*.

Parameters:
      * input (Tensor) -- the input tensor.
  * **dims** (*a list** or **tuple*) -- axis to flip on

Example:
  >>> x = torch.arange(8).view(2, 2, 2)
  >>> x
  tensor([[[ 0,  1],
           [ 2,  3]],

          [[ 4,  5],
           [ 6,  7]]])
  >>> torch.flip(x, [0, 1])
  tensor([[[ 6,  7],
           [ 4,  5]],

          [[ 2,  3],
           [ 0,  1]]])
torch.median
torch.median(input) -> Tensor
Returns the median of the values in "input".
Note:
 The median is not unique for "input" tensors with an even number
 of elements. In this case the lower of the two medians is
 returned. To compute the mean of both medians, use
 "torch.quantile()" with "q=0.5" instead.

Warning:
 This function produces deterministic (sub)gradients unlike
 "median(dim=0)"

Parameters:
      input (Tensor) -- the input tensor.
Example:
  >>> a = torch.randn(1, 3)
  >>> a
  tensor([[ 1.5219, -1.5212,  0.2202]])
  >>> torch.median(a)
  tensor(0.2202)

torch.median(input, dim=- 1, keepdim=False, *, out=None)
Returns a namedtuple "(values, indices)" where "values" contains
   the median of each row of "input" in the dimension "dim", and
   "indices" contains the index of the median values found in the
   dimension "dim".
By default, "dim" is the last dimension of the "input" tensor.
If "keepdim" is "True", the output tensors are of the same size as
   "input" except in the dimension "dim" where they are of size 1.
   Otherwise, "dim" is squeezed (see "torch.squeeze()"), resulting in
   the outputs tensor having 1 fewer dimension than "input".
Note:
 The median is not unique for "input" tensors with an even number
 of elements in the dimension "dim". In this case the lower of the
 two medians is returned. To compute the mean of both medians in
 "input", use "torch.quantile()" with "q=0.5" instead.

Warning:
 "indices" does not necessarily contain the first occurrence of
 each median value found, unless it is unique. The exact
 implementation details are device-specific. Do not expect the
 same result when run on CPU and GPU in general. For the same
 reason do not expect the gradients to be deterministic.

Parameters:
      * input (Tensor) -- the input tensor.
  * **dim** (*int*) -- the dimension to reduce.
keepdim (bool) -- whether the output tensor has "dim"
        retained or not.

Keyword Arguments:
      out ((Tensor, Tensor), optional) -- The first
      tensor will be populated with the median values and the second
      tensor, which must have dtype long, with their indices in the
      dimension "dim" of "input".
Example:
  >>> a = torch.randn(4, 5)
  >>> a
  tensor([[ 0.2505, -0.3982, -0.9948,  0.3518, -1.3131],
          [ 0.3180, -0.6993,  1.0436,  0.0438,  0.2270],
          [-0.2751,  0.7303,  0.2192,  0.3321,  0.2488],
          [ 1.0778, -1.9510,  0.7048,  0.4742, -0.7125]])
  >>> torch.median(a, 1)
  torch.return_types.median(values=tensor([-0.3982,  0.2270,  0.2488,  0.4742]), indices=tensor([1, 4, 4, 3]))
ConstantPad1d
class torch.nn.ConstantPad1d(padding, value)
Pads the input tensor boundaries with a constant value.
For N-dimensional padding, use "torch.nn.functional.pad()".
Parameters:
      padding (int, tuple) -- the size of the padding. If is
      int, uses the same padding in both boundaries. If a 2-tuple,
      uses (\text{padding_left}, \text{padding_right})
Shape:
      * Input: (C, W_{in}) or (N, C, W_{in}).
  * Output: (C, W_{out}) or (N, C, W_{out}), where

    W_{out} = W_{in} + \text{padding\_left} +
    \text{padding\_right}

Examples:
  >>> m = nn.ConstantPad1d(2, 3.5)
  >>> input = torch.randn(1, 2, 4)
  >>> input
  tensor([[[-1.0491, -0.7152, -0.0749,  0.8530],
           [-1.3287,  1.8966,  0.1466, -0.2771]]])
  >>> m(input)
  tensor([[[ 3.5000,  3.5000, -1.0491, -0.7152, -0.0749,  0.8530,  3.5000,
             3.5000],
3.5000],
               [ 3.5000,  3.5000, -1.3287,  1.8966,  0.1466, -0.2771,  3.5000,
                 3.5000]]])
      >>> m = nn.ConstantPad1d(2, 3.5)
      >>> input = torch.randn(1, 2, 3)
      >>> input
      tensor([[[ 1.6616,  1.4523, -1.1255],
               [-3.6372,  0.1182, -1.8652]]])
      >>> m(input)
      tensor([[[ 3.5000,  3.5000,  1.6616,  1.4523, -1.1255,  3.5000,  3.5000],
               [ 3.5000,  3.5000, -3.6372,  0.1182, -1.8652,  3.5000,  3.5000]]])
      >>> # using different paddings for different sides
      >>> m = nn.ConstantPad1d((3, 1), 3.5)
      >>> m(input)
      tensor([[[ 3.5000,  3.5000,  3.5000,  1.6616,  1.4523, -1.1255,  3.5000],
               [ 3.5000,  3.5000,  3.5000, -3.6372,  0.1182, -1.8652,  3.5000]]])
ZeroPad2d
class torch.nn.ZeroPad2d(padding)
Pads the input tensor boundaries with zero.
For N-dimensional padding, use "torch.nn.functional.pad()".
Parameters:
      padding (int, tuple) -- the size of the padding. If is
      int, uses the same padding in all boundaries. If a 4-tuple,
      uses (\text{padding_left}, \text{padding_right},
      \text{padding_top}, \text{padding_bottom})
Shape:
      * Input: (N, C, H_{in}, W_{in}) or (C, H_{in}, W_{in}).
  * Output: (N, C, H_{out}, W_{out}) or (C, H_{out}, W_{out}),
    where

    H_{out} = H_{in} + \text{padding\_top} +
    \text{padding\_bottom}

    W_{out} = W_{in} + \text{padding\_left} +
    \text{padding\_right}

Examples:
  >>> m = nn.ZeroPad2d(2)
  >>> input = torch.randn(1, 1, 3, 3)
  >>> input
  tensor([[[[-0.1678, -0.4418,  1.9466],
            [ 0.9604, -0.4219, -0.5241],
            [-0.9162, -0.5436, -0.6446]]]])
  >>> m(input)
m(input)
      tensor([[[[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
                [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
                [ 0.0000,  0.0000, -0.1678, -0.4418,  1.9466,  0.0000,  0.0000],
                [ 0.0000,  0.0000,  0.9604, -0.4219, -0.5241,  0.0000,  0.0000],
                [ 0.0000,  0.0000, -0.9162, -0.5436, -0.6446,  0.0000,  0.0000],
                [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
                [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]]])
      >>> # using different paddings for different sides
      >>> m = nn.ZeroPad2d((1, 1, 2, 0))
      >>> m(input)
      tensor([[[[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
                [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
                [ 0.0000, -0.1678, -0.4418,  1.9466,  0.0000],
                [ 0.0000,  0.9604, -0.4219, -0.5241,  0.0000],
                [ 0.0000, -0.9162, -0.5436, -0.6446,  0.0000]]]])
torch.linalg.ldl_solve
torch.linalg.ldl_solve(LD, pivots, B, *, hermitian=False, out=None) -> Tensor
Computes the solution of a system of linear equations using the LDL
   factorization.
"LD" and "pivots" are the compact representation of the LDL
   factorization and are expected to be computed by
   "torch.linalg.ldl_factor_ex()". "hermitian" argument to this
   function should be the same as the corresponding arguments in
   "torch.linalg.ldl_factor_ex()".
Supports input of float, double, cfloat and cdouble dtypes. Also
   supports batches of matrices, and if "A" is a batch of matrices
   then the output has the same batch dimensions.
Warning:
 This function is "experimental" and it may change in a future
 PyTorch release.

Parameters:
      * LD (Tensor) -- the n times n matrix or the batch of
        such matrices of size (, n, n)* where *** is one or more
        batch dimensions.
batch dimensions.
  * **pivots** (*Tensor*) -- the pivots corresponding to the LDL
    factorization of "LD".

  * **B** (*Tensor*) -- right-hand side tensor of shape *(*, n,
    k)*.

Keyword Arguments:
      * hermitian (bool, optional) -- whether to consider
        the decomposed matrix to be Hermitian or symmetric. For real-
        valued matrices, this switch has no effect. Default: False.
  * **out** (*tuple**, **optional*) -- output tensor. *B* may be
    passed as *out* and the result is computed in-place on *B*.
    Ignored if *None*. Default: *None*.

Examples:
  >>> A = torch.randn(2, 3, 3)
  >>> A = A @ A.mT # make symmetric
  >>> LD, pivots, info = torch.linalg.ldl_factor_ex(A)
  >>> B = torch.randn(2, 3, 4)
  >>> X = torch.linalg.ldl_solve(LD, pivots, B)
  >>> torch.linalg.norm(A @ X - B)
  >>> tensor(0.0001)
torch.tan
torch.tan(input, *, out=None) -> Tensor
Returns a new tensor with the tangent of the elements of "input".
  \text{out}_{i} = \tan(\text{input}_{i})

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> a = torch.randn(4)
  >>> a
  tensor([-1.2027, -1.7687,  0.4412, -1.3856])
  >>> torch.tan(a)
  tensor([-2.5930,  4.9859,  0.4722, -5.3366])
leaky_relu
class torch.ao.nn.quantized.functional.leaky_relu(input, negative_slope=0.01, inplace=False, scale=None, zero_point=None)
Quantized version of the. leaky_relu(input, negative_slope=0.01,
   inplace=False, scale, zero_point) -> Tensor
Applies element-wise, \text{LeakyReLU}(x) = \max(0, x) +
   \text{negative_slope} * \min(0, x)
Parameters:
      * input (Tensor) -- Quantized input
  * **negative_slope** (*float*) -- The slope of the negative
    input

  * **inplace** (*bool*) -- Inplace modification of the input
    tensor

  * **scale** (*Optional**[**float**]*) -- Scale and zero point of
    the output tensor.

  * **zero_point** (*Optional**[**int**]*) -- Scale and zero point
    of the output tensor.

See "LeakyReLU" for more details.
torch.func.jacfwd
torch.func.jacfwd(func, argnums=0, has_aux=False, *, randomness='error')
Computes the Jacobian of "func" with respect to the arg(s) at index
   "argnum" using forward-mode autodiff
Parameters:
      * func (function) -- A Python function that takes one or
        more arguments, one of which must be a Tensor, and returns one
        or more Tensors
  * **argnums** (*int** or **Tuple**[**int**]*) -- Optional,
    integer or tuple of integers, saying which arguments to get
    the Jacobian with respect to. Default: 0.

  * **has_aux** (*bool*) -- Flag indicating that "func" returns a
    "(output, aux)" tuple where the first element is the output of
    the function to be differentiated and the second element is
    auxiliary objects that will not be differentiated. Default:
    False.

  * **randomness** (*str*) -- Flag indicating what type of
    randomness to use. See "vmap()" for more detail. Allowed:
"different", "same", "error". Default: "error"
Returns:
      Returns a function that takes in the same inputs as "func" and
      returns the Jacobian of "func" with respect to the arg(s) at
      "argnums". If "has_aux is True", then the returned function
      instead returns a "(jacobian, aux)" tuple where "jacobian" is
      the Jacobian and "aux" is auxiliary objects returned by "func".
Note:
 You may see this API error out with "forward-mode AD not
 implemented for operator X". If so, please file a bug report and
 we will prioritize it. An alternative is to use "jacrev()", which
 has better operator coverage.

A basic usage with a pointwise, unary operation will give a
   diagonal array as the Jacobian



from torch.func import jacfwd
x = torch.randn(5)
jacobian = jacfwd(torch.sin)(x)
expected = torch.diag(torch.cos(x))
assert torch.allclose(jacobian, expected)



"jacfwd()" can be composed with vmap to produce batched Jacobians: