mlx.core.trace

目录

mlx.core.trace#

trace(a: array, /, offset: int = 0, axis1: int = 0, axis2: int = 1, dtype: Dtype | None = None, *, stream: None | Stream | Device = None) array#

返回给定数组中指定对角线的元素之和。

参数:
  • a (array) – 输入数组

  • offset (int, 可选) – 对角线相对于主对角线的偏移量。可以是正数或负数。默认值: 0

  • axis1 (int, 可选) – 应从中提取对角线的 2-D 子数组的第一个轴。默认值: 0

  • axis2 (int, 可选) – 应从中提取对角线的 2-D 子数组的第二个轴。默认值: 1

  • dtype (Dtype, 可选) – 输出数组的数据类型。如果未指定,则输出类型从输入数组推断。

返回值:

指定对角线的元素之和。

返回类型:

array