torch.func.linearize¶
- torch.func.linearize(func, *primals)[source]¶
返回
func
在primals
处的值以及在primals
处的线性近似。- 参数
func (Callable) – 一个接受一个或多个参数的 Python 函数。
primals (Tensors) –
func
的位置参数,必须全部为 Tensor。这些是函数线性近似的值。
- 返回
返回一个
(output, jvp_fn)
元组,其中包含应用于primals
的func
的输出,以及一个计算在primals
处评估的func
的 jvp 的函数。- 返回类型
如果要在
primals
处多次计算 jvp,则 linearize 非常有用。但是,为了实现这一点,linearize 会保存中间计算,并且比直接应用 jvp 具有更高的内存要求。因此,如果所有tangents
都是已知的,则使用 vmap(jvp) 而不是 linearize 可能会更有效。注意
linearize 评估
func
两次。请提交 issue 以获取单次评估的实现。- 示例:
>>> import torch >>> from torch.func import linearize >>> def fn(x): ... return x.sin() ... >>> output, jvp_fn = linearize(fn, torch.zeros(3, 3)) >>> jvp_fn(torch.ones(3, 3)) tensor([[1., 1., 1.], [1., 1., 1.], [1., 1., 1.]]) >>>