JAX 中的矢量化最小化和寻根-解网

问：

我有一系列参数化的函数args

f(x, args)

并希望确定的值的最小值。我可以访问该函数及其衍生物。我的第一次尝试是循环遍历的不同值并在每次迭代中使用 scipy.optimizer，但这需要很长时间。我相信可以通过矢量化来加快操作速度。我的下一个尝试是在 or 中使用，但我似乎无法为传递多个值。fxN = 1000argsargsjax.vmapjax.scipy.optimize.minimizejaxopt.ScipyMinimizeargs

或者，我可以编写自己的矢量化优化方法，例如平分，其中矢量化是指对数组进行固定次数的迭代操作，如果其中一个优化问题提前达到一定的容错水平，则不会提前停止。我希望使用一些优化的现成算法。

如果 jax 中有一个实现可用，我希望使用一些已经优化的现成算法。这个线程是相关的，但没有改变。args

蟒优化西皮矢量化贾克斯

import jax
import jax.numpy as jnp
from jax.scipy import optimize

def f(x, args):
  a, b = args
  return jnp.sum(a + (x - b) ** 2)

def find_min(a, b):
  x0 = jnp.array([1.0])
  args = (a, b)
  return optimize.minimize(f, x0, (args,), method="BFGS")

a_grid, b_grid = jnp.meshgrid(jnp.arange(5.0), jnp.arange(5.0))

results = jax.vmap(find_min)(a_grid.ravel(), b_grid.ravel())

print(results.success)
# [ True  True  True  True  True  True  True  True  True  True  True  True
#   True  True  True  True  True  True  True  True  True  True  True  True
#   True]

print(results.x.T)
# [[0. 0. 0. 0. 0. 1. 1. 1. 1. 1. 2. 2. 2. 2. 2.
#   3. 3. 3. 3. 3. 4. 4. 4. 4. 4.]]

上一个：使用列名列表在 pandas 数据帧中执行操作

下一个：从一长串字典中高效提取元素

JAX 中的矢量化最小化和寻根

vectorized minimization and root finding in jax

评论