scico.function#
Function class.
Classes
|
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 multiplearray-like
arguments to anotherarray-like
. It is more general than bothFunctional
, which is a mapping to a scalar, andOperator
, 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 toNone
. IfNone
, output_shape is determined by evaluating self.__call__ on input arrays of zeros.eval_fn (
Optional
[Callable
]) – Function used in evaluating thisFunction
. Defaults toNone
. 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 tofloat32
.output_dtype (
Optional
[DType
]) – dtype for output argument. Defaults toNone
. IfNone
, output_dtype is determined by evaluating self.__call__ on an input arrays of zeros.jit (
bool
) – IfTrue
, 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:
- Returns:
Value of function with specified parameters.
- join()[source]#
Combine inputs into a
BlockArray
.Construct an equivalent
Operator
taking a singleBlockArray
input combining all inputs of thisFunction
.- Return type:
- Returns:
An
Operator
taking aBlockArray
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 tojax.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:
- Return type:
- 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 theOperator
should be included (as the first component of aBlockArray
) in the output of the JacobianLinearOperator
constructed by this function.
- Return type:
- Returns:
A
LinearOperator
capable of computing Jacobian-vector products.