jax.lax.conv#

jax.lax.conv(lhs, rhs, window_strides, padding, precision=None, preferred_element_type=None)[source]#

conv_general_dilated 的便捷封装。

参数:
  • lhs (Array) – 一个秩为 n+2 维的输入数组。

  • rhs (Array) – 一个秩为 n+2 维的内核权重数组。

  • window_strides (Sequence[int]) – 一个 n 个整数的序列,表示窗口间的步幅。

  • padding (str) – 字符串 ‘SAME’,或字符串 ‘VALID’

  • precision (lax.PrecisionLike | None) – 可选。可以是 None,表示后端的默认精度,一个 Precision 枚举值 (Precision.DEFAULT, Precision.HIGHPrecision.HIGHEST),或一对 Precision 枚举值,表示 lhs`rhs 的精度。

  • preferred_element_type (DTypeLike | None | None) – 可选。可以是 None,表示输入类型的默认累积类型,或一个数据类型,表示将结果累积到该数据类型并返回该数据类型的结果。

返回:

包含卷积结果的数组。

返回类型:

Array