scico.function

Function class.

Classes

Function(input_shapes[, output_shape, ...])

Function class.
class scico.function.Function(input_shapes, output_shape=None, eval_fn=None, input_dtypes=<class 'jax.numpy.float32'>, output_dtype=None, jit=False)[source]

Bases: object

Function class.

A Function maps multiple array-like arguments to another array-like. It is more general than both Functional, which is a mapping to a scalar, and Operator, which takes a single argument.

Parameters:

  • input_shapes (Sequence[Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...]]]) – Shapes of input arrays.

  • output_shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...], ...], None]) – Shape of output array. Defaults to None. If None, output_shape is determined by evaluating self.__call__ on input arrays of zeros.

  • eval_fn (Optional[Callable]) – Function used in evaluating this Function. Defaults to None. Required unless __init__ is being called from a derived class with an _eval method.

  • input_dtypes (Union[DType, Sequence[DType]]) – dtype for input argument. If a single dtype is specified, it implies a common dtype for all inputs, otherwise a list or tuple of values should be provided, one per input. Defaults to float32.

  • output_dtype (Optional[DType]) – dtype for output argument. Defaults to None. If None, output_dtype is determined by evaluating self.__call__ on an input arrays of zeros.

  • jit (bool) – If True, jit the evaluation function.

__call__(*args)[source]

Evaluate this function with the specified parameters.

Parameters:

*args (Union[Array, BlockArray]) – Parameters at which to evaluate the function.

Return type:

Union[Array, BlockArray]

Returns:

Value of function with specified parameters.

slice(index, *fix_args)[source]

Fix all but one parameter, returning a Operator.

Parameters:

  • index (int) – Index of parameter that remains free.

  • *fix_args (Union[Array, BlockArray]) – Fixed values for remaining parameters.

Return type:

Operator

Returns:

An Operator taking the free parameter of the Function as its input.

join()[source]

Combine inputs into a BlockArray.

Construct an equivalent Operator taking a single BlockArray input combining all inputs of this Function.

Return type:

Operator

Returns:

An Operator taking a BlockArray as its input.

jvp(index, v, *args)[source]

Jacobian-vector product with respect to a single parameter.

Compute a Jacobian-vector product with respect to a single parameter of a Function. Note that the order of the parameters specifying where to evaluate the Jacobian and the vector in the product is reverse with respect to jax.jvp.

Parameters:

  • index (int) – Index of parameter with respect to which the Jacobian is to be computed.

  • v (Union[Array, BlockArray]) – Vector against which the Jacobian-vector product is to be computed.

  • *args (Union[Array, BlockArray]) – Values of function parameters at which Jacobian is to be computed.

Return type:

Tuple[Union[Array, BlockArray], Union[Array, BlockArray]]

Returns:

A pair consisting of the operator evaluated at the parameters specified by *args and the Jacobian-vector product.

vjp(index, *args, conjugate=True)[source]

Vector-Jacobian product with respect to a single parameter.

Compute a vector-Jacobian product with respect to a single parameter of a Function.

Parameters:

  • index (int) – Index of parameter with respect to which the Jacobian is to be computed.

  • *args (Union[Array, BlockArray]) – Values of function parameters at which Jacobian is to be computed.

  • conjugate (Optional[bool]) – If True, compute the product using the conjugate (Hermitian) transpose.

Return type:

Tuple[Tuple[Any, ...], Callable]

Returns:

A pair consisting of the operator evaluated at the parameters specified by *args and a function that computes the vector-Jacobian product.

jacobian(index, *args, include_eval=False)[source]

Construct Jacobian linear operator for the function.

Construct a Jacobian LinearOperator that computes vector products with the Jacobian with respect to a specified variable of the function.

Parameters:

  • index (int) – Index of parameter with respect to which the Jacobian is to be computed.

  • *args (Union[Array, BlockArray]) – Values of function parameters at which Jacobian is to be computed.

  • include_eval (Optional[bool]) – Flag indicating whether the result of evaluating the Operator should be included (as the first component of a BlockArray) in the output of the Jacobian LinearOperator constructed by this function.

Return type:

LinearOperator

Returns:

A LinearOperator capable of computing Jacobian-vector products.