mlx_rs::transforms

Function vjp

Source 
pub fn vjp<'a, F>(
    f: F,
    primals: &[Array],
    cotangents: &[Array],
) -> Result<(Vec<Array>, Vec<Array>)>
where
    F: FnMut(&[Array]) -> Vec<Array> + 'a,
Expand description

Compute the vector-Jacobian product.

Computes the product of the cotangents with the Jacobian of a function f evaluated at primals.

§Params:

  • f: function which takes an array of Array and returns an array of Array
  • primals: array of Array at which to evaluate the Jacobian
  • cotangents: array of Array which are the “vector” in the vector-Jacobian product. The cotangents should be the same in number, shape and type as the outputs of f

§Returns:

array of the vector-Jacobian products which is the same in number, shape and type of the outputs of f