jax.numpy.divmod

jax.numpy.divmod(x1, x2, /)

逐元素计算 x1 除以 x2 的整数商和余数

numpy.divmod 的 JAX 实现。

参数:

• x1 (类数组) – 输入数组，被除数

• x2 (类数组) – 输入数组，除数

返回:

数组元组 (x1 // x2, x1 % x2)。

返回类型:

tuple[Array, Array]

另请参阅

• jax.numpy.floor_divide()：向下取整除法函数

• jax.numpy.remainder(): 求余函数

示例

>>> x1 = jnp.array([10, 20, 30])
>>> x2 = jnp.array([3, 4, 7])
>>> jnp.divmod(x1, x2)
(Array([3, 5, 4], dtype=int32), Array([1, 0, 2], dtype=int32))

>>> x1 = jnp.array([-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5])
>>> x2 = 3
>>> jnp.divmod(x1, x2)
(Array([-2, -2, -1, -1, -1,  0,  0,  0,  1,  1,  1], dtype=int32),
 Array([1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2], dtype=int32))

>>> x1 = jnp.array([6, 6, 6], dtype=jnp.int32)
>>> x2 = jnp.array([1.9, 2.5, 3.1], dtype=jnp.float32)
>>> jnp.divmod(x1, x2)
(Array([3., 2., 1.], dtype=float32),
 Array([0.30000007, 1.        , 2.9       ], dtype=float32))