jax.numpy.inner#
- jax.numpy.inner(a, b, *, precision=None, preferred_element_type=None)[source]#
计算两个数组的内积。
JAX 实现的
numpy.inner()
。与
jax.numpy.matmul()
或jax.numpy.dot()
不同,这始终沿着每个输入的最后一个维度执行收缩。- 参数:
a (Array | ndarray | bool | number | bool | int | float | complex) – 形状为
(..., N)
的数组b (Array | ndarray | bool | number | bool | int | float | complex) – 形状为
(..., N)
的数组precision (None | str | Precision | tuple[str, str] | tuple[Precision, Precision] | DotAlgorithm | DotAlgorithmPreset) –
None
(默认),表示后端的默认精度,Precision
枚举值 (Precision.DEFAULT
,Precision.HIGH
或Precision.HIGHEST
) 或两个此类值的元组,指示a
和b
的精度。preferred_element_type (str | type[Any] | dtype | SupportsDType | None) –
None
(默认),表示输入类型的默认累积类型,或数据类型,指示累积结果并返回具有该数据类型的结果。
- 返回:
形状为
(*a.shape[:-1], *b.shape[:-1])
的数组,包含输入的批量向量积。- 返回类型:
另请参阅
jax.numpy.vecdot()
: 沿指定轴的共轭乘法。jax.numpy.tensordot()
: 一般张量乘法。jax.numpy.matmul()
: 一般批量矩阵和向量乘法。
示例
对于 1D 输入,这实现了标准(非共轭)向量乘法
>>> a = jnp.array([1j, 3j, 4j]) >>> b = jnp.array([4., 2., 5.]) >>> jnp.inner(a, b) Array(0.+30.j, dtype=complex64)
对于多维输入,批量维度堆叠而不是广播
>>> a = jnp.ones((2, 3)) >>> b = jnp.ones((5, 3)) >>> jnp.inner(a, b).shape (2, 5)