Tensorflow.js tf.sum()函数

在TensorFlow.js中，tf.sum()函数用于计算张量的所有元素的和。

语法

以下是tf.sum()函数的语法：

tf.sum(inputTensor, axis, keepDims)

其中：

• inputTensor：一个张量，用于计算。
• axis：一个可选参数，指定要沿着哪个轴计算和。默认情况下，计算所有维度的和。可以传递一个整数数组，以指定要沿着哪些轴计算和。
• keepDims：一个可选参数，指定是否保持输出张量的shape与输入张量的shape相同。

返回值

tf.sum()函数返回所有元素的和。

示例

以下是tf.sum()函数的一个简单示例：

const inputTensor = tf.tensor2d([[1, 2], [3, 4]]);
const sum = tf.sum(inputTensor);
sum.print(); // 输出: 10

在这个例子中，输入张量是一个二维张量，tf.sum()函数将计算它的所有元素的和，并返回结果10。

沿着指定轴计算和

我们可以将axis参数传递给tf.sum()函数，以指定要沿着哪个轴计算和。以下是一个示例：

const inputTensor = tf.tensor2d([[1, 2], [3, 4]]);
const sum = tf.sum(inputTensor, 0);
sum.print(); // 输出: [4, 6]

在这个例子中，tf.sum()函数沿着第一个维度计算和，返回每列的和。

保持输出张量的shape

我们可以将keepDims参数传递给tf.sum()函数，以保持输出张量的shape与输入张量的shape相同。以下是一个示例：

const inputTensor = tf.tensor2d([[1, 2], [3, 4]]);
const sum = tf.sum(inputTensor, 0, true);
sum.print(); // 输出: [[4, 6]]

在这个例子中，tf.sum()函数沿着第一个维度计算和，并返回一个二维张量，shape为[1,2]。

总结

tf.sum()函数是TensorFlow.js中一个用于计算输入张量所有元素的和的函数。我们可以用它来计算整个张量的和，或者沿着指定轴计算和。我们还可以使用keepDims参数来保持输出张量的shape与输入张量的shape相同。