auxiliary objects that will not be differentiated. Default:
        False.
  * **chunk_size** (*None** or **int*) -- If None (default), use
    the maximum chunk size (equivalent to doing a single vmap over
    vjp to compute the jacobian). If 1, then compute the jacobian
    row-by-row with a for-loop. If not None, then compute the
    jacobian "chunk_size" rows at a time (equivalent to doing
    multiple vmap over vjp). If you run into memory issues
    computing the jacobian, please try to specify a non-None
    chunk_size.

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".
A basic usage with a pointwise, unary operation will give a
   diagonal array as the Jacobian