jax.numpy.roots

内容

jax.numpy.roots#

jax.numpy.roots(p, *, strip_zeros=True)[source]#

返回给定系数 p 的多项式的根。

JAX 实现的 numpy.roots()。

参数：

  • p (ArrayLike) – 秩为 1 的多项式系数数组。

  • strip_zeros (bool) – 布尔值，默认为 True。如果为 True，则系数中的前导零将被去除，类似于 numpy.roots()。如果设置为 False，则不会去除前导零，并且未定义的根将在函数输出中以 NaN 值表示。strip_zeros 必须设置为 False，才能使函数与 jax.jit() 和其他 JAX 变换兼容。

返回:

包含多项式根的数组。

返回类型:

数组

注意

与该函数的 np.roots 不同，jnp.roots 无论根的值如何，都以复数数组的形式返回根。

另请参阅

示例

>>> coeffs = jnp.array([0, 1, 2])

默认行为与 numpy 一致，并去除前导零

>>> jnp.roots(coeffs)
Array([-2.+0.j], dtype=complex64)

使用 strip_zeros=False，额外的根将设置为 NaN

>>> jnp.roots(coeffs, strip_zeros=False)
Array([-2. +0.j, nan+nanj], dtype=complex64)
内容