API 参考

本节提供了 Riemann 库中所有函数、类和模块的全面参考。

张量操作

张量创建函数

riemann.tensor(data, dtype=None, requires_grad=False)

从数据创建张量。

参数:

data (array_like) – 初始化张量的数据
dtype (numpy.dtype, optional) – 张量的期望数据类型
requires_grad (bool, optional) – 是否跟踪此张量上的操作

返回:

包含给定数据的张量

返回类型:

riemann.TN

riemann.zeros(*shape, dtype=None, requires_grad=False)

创建一个填充零的张量。

参数:

shape (int 或整数元组) – 张量的形状
dtype (numpy.dtype, optional) – 张量的期望数据类型
requires_grad (bool, optional) – 是否跟踪此张量上的操作

返回:

填充零的张量

返回类型:

riemann.TN

riemann.ones(*shape, dtype=None, requires_grad=False)

创建一个填充一的张量。

参数:

shape (int 或整数元组) – 张量的形状
dtype (numpy.dtype, optional) – 张量的期望数据类型
requires_grad (bool, optional) – 是否跟踪此张量上的操作

返回:

填充一的张量

返回类型:

riemann.TN

riemann.empty(*shape, dtype=None, requires_grad=False)

创建一个未初始化的张量。

参数:

shape (int 或整数元组) – 张量的形状
dtype (numpy.dtype, optional) – 张量的期望数据类型
requires_grad (bool, optional) – 是否跟踪此张量上的操作

返回:

未初始化的张量

返回类型:

riemann.TN

riemann.full(*shape, fill_value, dtype=None, requires_grad=False)

创建一个填充特定值的张量。

参数:

shape (int 或整数元组) – 张量的形状
fill_value (标量) – 填充张量的值
dtype (numpy.dtype, optional) – 张量的期望数据类型
requires_grad (bool, optional) – 是否跟踪此张量上的操作

返回:

填充指定值的张量

返回类型:

riemann.TN

riemann.eye(n, m=None, dtype=None, requires_grad=False)

创建一个对角线为一，其余为零的二维张量。

参数:

n (int) – 行数
m (int, optional) – 列数（默认为 n）
dtype (numpy.dtype, optional) – 张量的期望数据类型
requires_grad (bool, optional) – 是否跟踪此张量上的操作

返回:

对角线为一的二维张量

返回类型:

riemann.TN

riemann.zeros_like(tsr, dtype=None, requires_grad=False)

创建一个与输入张量形状相同的零张量。

参数:

tsr (riemann.TN) – 参考张量
dtype (numpy.dtype, optional) – 张量的期望数据类型
requires_grad (bool, optional) – 是否跟踪此张量上的操作

返回:

与输入张量形状相同的零张量

返回类型:

riemann.TN

riemann.ones_like(tsr, dtype=None, requires_grad=False)

创建一个与输入张量形状相同的一张量。

参数:

tsr (riemann.TN) – 参考张量
dtype (numpy.dtype, optional) – 张量的期望数据类型
requires_grad (bool, optional) – 是否跟踪此张量上的操作

返回:

与输入张量形状相同的一张量

返回类型:

riemann.TN

riemann.empty_like(tsr, dtype=None, requires_grad=False)

创建一个与输入张量形状相同的未初始化张量。

参数:

tsr (riemann.TN) – 参考张量
dtype (numpy.dtype, optional) – 张量的期望数据类型
requires_grad (bool, optional) – 是否跟踪此张量上的操作

返回:

与输入张量形状相同的未初始化张量

返回类型:

riemann.TN

riemann.full_like(tsr, fill_value, dtype=None, requires_grad=False)

创建一个与输入张量形状相同的填充特定值的张量。

参数:

tsr (riemann.TN) – 参考张量
fill_value (标量) – 填充张量的值
dtype (numpy.dtype, optional) – 张量的期望数据类型
requires_grad (bool, optional) – 是否跟踪此张量上的操作

返回:

与输入张量形状相同的填充指定值的张量

返回类型:

riemann.TN

随机数生成

riemann.rand(*size, requires_grad=False, dtype=None)

创建一个填充来自 [0, 1) 均匀分布的随机数的张量。

参数:

size (int 或整数元组) – 张量的形状
requires_grad (bool, optional) – 是否跟踪此张量上的操作
dtype (numpy.dtype, optional) – 张量的期望数据类型

返回:

填充随机值的张量

返回类型:

riemann.TN

riemann.randn(*size, requires_grad=False, dtype=None)

创建一个填充来自标准正态分布的随机数的张量。

参数:

size (int 或整数元组) – 张量的形状
requires_grad (bool, optional) – 是否跟踪此张量上的操作
dtype (numpy.dtype, optional) – 张量的期望数据类型

返回:

填充随机值的张量

返回类型:

riemann.TN

riemann.randint(low, high, size, requires_grad=False, dtype=int64)

创建一个填充从 low（包含）到 high（不包含）的随机整数的张量。

参数:

low (int) – 要抽取的最小整数
high (int) – 要抽取的最大整数加一
size (int 或整数元组) – 张量的形状
requires_grad (bool, optional) – 是否跟踪此张量上的操作
dtype (numpy.dtype, optional) – 张量的期望数据类型

返回:

填充随机整数的张量

返回类型:

riemann.TN

riemann.randperm(n, requires_grad=False, dtype=int64)

创建一个包含 0 到 n-1 随机顺序数字的张量。

参数:

n (int) – 上限（不包含）
requires_grad (bool, optional) – 是否跟踪此张量上的操作
dtype (numpy.dtype, optional) – 张量的期望数据类型

返回:

包含随机排列整数的张量

返回类型:

riemann.TN

riemann.normal(mean, std, size, dtype=None)

创建一个填充来自正态分布的随机数的张量。

参数:

mean (float) – 正态分布的均值
std (float) – 正态分布的标准差
size (int 或整数元组) – 张量的形状
dtype (numpy.dtype, optional) – 张量的期望数据类型

返回:

填充随机值的张量

返回类型:

riemann.TN

随机种子控制

riemann.manual_seed(seed)

设置随机数生成器的种子，用于确保随机操作的可重复性。

参数:: seed (int) – 随机种子值
返回:: 随机数生成器对象
返回类型:: torch.Generator

序列和范围函数

riemann.arange(start, end=None, step=1.0, dtype=None, requires_grad=False)

创建一个从 start 到 end 以 step 为步长的 1-D 张量。

参数:

start (float) – 起始值
end (float, optional) – 结束值（不包含）
step (float, optional) – 值之间的间距
dtype (numpy.dtype, optional) – 张量的期望数据类型
requires_grad (bool, optional) – 是否跟踪此张量上的操作

返回:

包含等间距值的 1-D 张量

返回类型:

riemann.TN

riemann.linspace(start, end, steps=100, endpoint=True, dtype=None, requires_grad=False)

创建一个在给定区间内包含等间距值的 1-D 张量。

参数:

start (float) – 起始值
end (float) – 结束值
steps (int, optional) – 要生成的样本数
endpoint (bool, optional) – 是否包含结束值
dtype (numpy.dtype, optional) – 张量的期望数据类型
requires_grad (bool, optional) – 是否跟踪此张量上的操作

返回:

包含等间距值的 1-D 张量

返回类型:

riemann.TN

张量属性

riemann.TN.dtype()

返回张量的数据类型。

返回:: 张量的数据类型
返回类型:: numpy.dtype

riemann.TN.real()

返回复数张量的实部。

返回:: 包含实部的张量
返回类型:: riemann.TN

riemann.TN.imag()

返回复数张量的虚部。

返回:: 包含虚部的张量
返回类型:: riemann.TN

riemann.TN.shape()

返回张量的形状。

返回:: 张量各维度大小的元组
返回类型:: tuple

riemann.TN.ndim()

返回张量的维度数。

返回:: 张量的维度数
返回类型:: int

riemann.TN.device()

返回张量所在的设备。

返回:: 张量所在的设备对象
返回类型:: Device

riemann.TN.is_cuda()

检查张量是否在CUDA设备上。

返回:: 如果张量在CUDA设备上返回True，否则返回False
返回类型:: bool

riemann.TN.is_cpu()

检查张量是否在CPU设备上。

返回:: 如果张量在CPU设备上返回True，否则返回False
返回类型:: bool

riemann.TN.is_leaf()

检查张量是否为叶子节点。

返回:: 如果张量为叶子节点返回True，否则返回False
返回类型:: bool

riemann.TN.is_floating_point()

检查张量是否为浮点类型。

返回:: 如果张量是浮点类型返回True，否则返回False
返回类型:: bool

riemann.TN.is_complex()

检查张量是否为复数类型。

返回:: 如果张量是复数类型返回True，否则返回False
返回类型:: bool

riemann.TN.isreal()

判断张量元素是否为实数。

返回:: 布尔张量，True表示对应元素为实数
返回类型:: riemann.TN

riemann.TN.isinf()

判断张量元素是否为无穷大。

返回:: 布尔张量，True表示对应元素为无穷大
返回类型:: riemann.TN

riemann.TN.isnan()

判断张量元素是否为NaN（非数值）。

返回:: 布尔张量，True表示对应元素为NaN
返回类型:: riemann.TN

riemann.TN.conj()

返回张量的复数共轭。

返回:: 包含共轭元素的张量
返回类型:: riemann.TN

riemann.TN.size(dim=None)

返回张量的大小。

参数:: dim (int, optional) – 要查询的维度索引，如果为None返回整个形状
返回:: 如果dim为None返回形状元组，否则返回指定维度的大小
返回类型:: tuple or int

riemann.TN.numel()

返回张量中元素的总数。

返回:: 张量中元素的数量
返回类型:: int

riemann.TN.is_contiguous()

检查张量的内存布局是否是连续的。

返回:: 如果张量是连续的返回True，否则返回False
返回类型:: bool

张量形状操作

riemann.reshape(input, shape)

返回一个具有相同数据但不同形状的张量。

参数:

input (riemann.TN) – 输入张量
shape (整数元组) – 新形状

返回:

具有新形状的张量

返回类型:

riemann.TN

riemann.squeeze(input, dim=None)

从张量的形状中移除大小为 1 的维度。

参数:

input (riemann.TN) – 输入张量
dim (int, optional) – 要压缩的维度

返回:

压缩维度后的张量

返回类型:

riemann.TN

riemann.unsqueeze(input, dim)

在指定位置插入大小为 1 的维度。

参数:

input (riemann.TN) – 输入张量
dim (int) – 要扩展的维度

返回:

扩展维度后的张量

返回类型:

riemann.TN

riemann.transpose(input, dim0, dim1)

交换张量的两个维度。

参数:

input (riemann.TN) – 输入张量
dim0 (int) – 要交换的第一个维度
dim1 (int) – 要交换的第二个维度

返回:

交换维度后的张量

返回类型:

riemann.TN

riemann.TN.mT

矩阵转置，即张量最后两个维度间的转置。

对于行向量 (1, n)，转置为列向量 (n, 1)；对于列向量 (n, 1)，转置为行向量 (1, n)。

返回:: 转置后的张量
返回类型:: riemann.TN

riemann.is_contiguous(input)

检查张量的内存是否连续。

参数:: input (riemann.TN) – 输入张量
返回:: 内存是否连续
返回类型:: bool

riemann.contiguous(input)

返回一个内存连续的张量。

参数:: input (riemann.TN) – 输入张量
返回:: 内存连续的张量
返回类型:: riemann.TN

riemann.gather(input, dim, index)

沿指定维度收集张量中的元素。

参数:

input (riemann.TN) – 输入张量
dim (int) – 收集的维度
index (riemann.TN) – 索引张量

返回:

收集后的张量

返回类型:

riemann.TN

riemann.scatter(input, dim, index, src)

沿指定维度将源张量中的元素分散到目标张量中。

参数:

input (riemann.TN) – 目标张量
dim (int) – 分散的维度
index (riemann.TN) – 索引张量
src (riemann.TN) – 源张量

返回:

分散后的张量

返回类型:

riemann.TN

riemann.broadcast_to(input, size)

将张量广播到新形状。

参数:

input (riemann.TN) – 输入张量
size (整数元组) – 目标形状

返回:

广播后的张量

返回类型:

riemann.TN

riemann.flip(input, dims)

沿指定维度反转元素的顺序。

参数:

input (riemann.TN) – 输入张量
dims (int 的列表或元组) – 要翻转的维度

返回:

翻转后的张量

返回类型:

riemann.TN

riemann.split(ts, split_indices, dim=0)

将张量拆分为多个子张量。

参数:

ts (riemann.TN) – 输入张量
split_indices (int 或整数列表) – 拆分的索引
dim (int, optional) – 沿其拆分的维度

返回:

张量列表

返回类型:

riemann.TN 的列表

riemann.stack(tensors, dim=0)

沿新维度堆叠张量。

参数:

tensors (riemann.TN 的序列) – 要堆叠的张量序列
dim (int, optional) – 要插入的维度

返回:

堆叠后的张量

返回类型:

riemann.TN

riemann.cat(tensors, dim=0)

沿现有维度连接张量。

参数:

tensors (riemann.TN 的序列) – 要连接的张量序列
dim (int, optional) – 沿其连接的维度

返回:

连接后的张量

返回类型:

riemann.TN

riemann.concatenate(tensors, dim=0)

沿现有维度连接张量。

参数:

tensors (riemann.TN 的序列) – 要连接的张量序列
dim (int, optional) – 沿其连接的维度

返回:

连接后的张量

返回类型:

riemann.TN

riemann.vstack(tensors)

垂直堆叠张量（按行）。

参数:: tensors (riemann.TN 的序列) – 要堆叠的张量序列
返回:: 垂直堆叠的张量
返回类型:: riemann.TN

riemann.hstack(tensors)

水平堆叠张量（按列）。

参数:: tensors (riemann.TN 的序列) – 要堆叠的张量序列
返回:: 水平堆叠的张量
返回类型:: riemann.TN

riemann.dstack(tensors)

深度堆叠张量（沿第2维）。

1D张量会先reshape为(1, N, 1)，2D张量会reshape为(M, N, 1)，然后沿第2维堆叠。

参数:: tensors (riemann.TN 的序列) – 要堆叠的张量序列
返回:: 深度堆叠的张量
返回类型:: riemann.TN

riemann.tensor_split(input, indices_or_sections, dim=0)

将张量沿指定维度分割为多个子张量。

当 indices_or_sections 为整数时，表示将张量分割成N段；当 indices_or_sections 为列表时，表示在指定的索引位置进行分割。

参数:

input (riemann.TN) – 输入张量
indices_or_sections (int 或 list[int]) – 分割参数（段数或索引列表）
dim (int, optional) – 沿其分割的维度，默认为0

返回:

分割后的张量元组

返回类型:

tuple[riemann.TN, …]

riemann.vsplit(input, indices_or_sections)

垂直分割张量（沿第0维）。

将张量沿第0维（垂直方向）分割为多个子张量。

参数:

input (riemann.TN) – 输入张量
indices_or_sections (int 或 list[int]) – 分割参数（段数或索引列表）

返回:

分割后的张量元组

返回类型:

tuple[riemann.TN, …]

riemann.hsplit(input, indices_or_sections)

水平分割张量（沿第1维）。

将张量沿第1维（水平方向）分割为多个子张量。

参数:

input (riemann.TN) – 输入张量
indices_or_sections (int 或 list[int]) – 分割参数（段数或索引列表）

返回:

分割后的张量元组

返回类型:

tuple[riemann.TN, …]

riemann.dsplit(input, indices_or_sections)

深度分割张量（沿第2维）。

将3D+张量沿第2维（深度方向）分割为多个子张量。

参数:

input (riemann.TN) – 输入张量（至少3维）
indices_or_sections (int 或 list[int]) – 分割参数（段数或索引列表）

返回:

分割后的张量元组

返回类型:

tuple[riemann.TN, …]

张量运算符

Riemann框架支持丰富的张量运算符，包括算术运算符、比较运算符、位运算符和原地运算符。这些运算符可以直接作用于张量对象，支持自动微分，并遵循Python的运算符优先级规则。

算术运算符

__add__(other)

张量加法运算符，等价于 +。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 另一个张量或标量值
返回:: 加法结果张量
返回类型:: riemann.TN

__radd__(other)

反向张量加法运算符，当左操作数不是张量时使用。

参数:: other (int 或 float 或 complex) – 非张量左操作数
返回:: 加法结果张量
返回类型:: riemann.TN

__sub__(other)

张量减法运算符，等价于 -。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 另一个张量或标量值
返回:: 减法结果张量
返回类型:: riemann.TN

__rsub__(other)

反向张量减法运算符，当左操作数不是张量时使用。

参数:: other (int 或 float 或 complex) – 非张量左操作数
返回:: 减法结果张量
返回类型:: riemann.TN

__mul__(other)

张量乘法运算符，等价于 *。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 另一个张量或标量值
返回:: 乘法结果张量
返回类型:: riemann.TN

__rmul__(other)

反向张量乘法运算符，当左操作数不是张量时使用。

参数:: other (int 或 float 或 complex) – 非张量左操作数
返回:: 乘法结果张量
返回类型:: riemann.TN

__matmul__(other)

张量矩阵乘法运算符，等价于 @。

参数:: other (riemann.TN) – 另一个张量
返回:: 矩阵乘法结果张量
返回类型:: riemann.TN
抛出:: RuntimeError – 如果任一操作数是标量

__rmatmul__(other)

反向张量矩阵乘法运算符，当左操作数不是张量时使用。

参数:: other (int 或 float 或 complex) – 非张量左操作数
返回:: 矩阵乘法结果张量
返回类型:: riemann.TN

__truediv__(other)

张量除法运算符，等价于 /。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 另一个张量或标量值
返回:: 除法结果张量
返回类型:: riemann.TN

__rtruediv__(other)

反向张量除法运算符，当左操作数不是张量时使用。

参数:: other (int 或 float 或 complex) – 非张量左操作数
返回:: 除法结果张量
返回类型:: riemann.TN

__pow__(other)

张量幂运算符，等价于 **。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 指数张量或标量值
返回:: 幂运算结果张量
返回类型:: riemann.TN

__rpow__(other)

反向张量幂运算符，当左操作数不是张量时使用。

参数:: other (int 或 float 或 complex) – 非张量左操作数
返回:: 幂运算结果张量
返回类型:: riemann.TN

__pos__()

张量正号运算符，等价于 +。

返回:: 原张量
返回类型:: riemann.TN

__neg__()

张量负号运算符，等价于 -。

返回:: 取负结果张量
返回类型:: riemann.TN

比较运算符

__lt__(other)

张量小于运算符，等价于 <。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 另一个张量或标量值
返回:: 布尔结果张量
返回类型:: riemann.TN
抛出:: TypeError – 如果other是None

__le__(other)

张量小于等于运算符，等价于 <=。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 另一个张量或标量值
返回:: 布尔结果张量
返回类型:: riemann.TN
抛出:: TypeError – 如果other是None

__gt__(other)

张量大于运算符，等价于 >。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 另一个张量或标量值
返回:: 布尔结果张量
返回类型:: riemann.TN
抛出:: TypeError – 如果other是None

__ge__(other)

张量大于等于运算符，等价于 >=。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 另一个张量或标量值
返回:: 布尔结果张量
返回类型:: riemann.TN
抛出:: TypeError – 如果other是None

__eq__(other)

张量等于运算符，等价于 ==。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 另一个张量或标量值
返回:: 布尔结果张量
返回类型:: riemann.TN

__ne__(other)

张量不等于运算符，等价于 !=。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 另一个张量或标量值
返回:: 布尔结果张量
返回类型:: riemann.TN

位运算符

__and__(other)

张量按位与运算符，等价于 &。

参数:: other (riemann.TN 或 int) – 另一个张量或标量值
返回:: 按位与结果张量
返回类型:: riemann.TN

__or__(other)

张量按位或运算符，等价于 |。

参数:: other (riemann.TN 或 int) – 另一个张量或标量值
返回:: 按位或结果张量
返回类型:: riemann.TN

__xor__(other)

张量按位异或运算符，等价于 ^。

参数:: other (riemann.TN 或 int) – 另一个张量或标量值
返回:: 按位异或结果张量
返回类型:: riemann.TN

__invert__()

张量按位取反运算符，等价于 ~。

返回:: 按位取反结果张量
返回类型:: riemann.TN

__lshift__(other)

张量左移位运算符，等价于 <<。

参数:: other (riemann.TN 或 int) – 移位位数
返回:: 左移位结果张量
返回类型:: riemann.TN

__rshift__(other)

张量右移位运算符，等价于 >>。

参数:: other (riemann.TN 或 int) – 移位位数
返回:: 右移位结果张量
返回类型:: riemann.TN

原地运算符

__iadd__(other)

张量原地加法运算符，等价于 +=。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 另一个张量或标量值
返回:: 原地修改后的张量
返回类型:: riemann.TN
抛出:: RuntimeError – 如果张量是需要梯度的叶子节点

__isub__(other)

张量原地减法运算符，等价于 -=。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 另一个张量或标量值
返回:: 原地修改后的张量
返回类型:: riemann.TN
抛出:: RuntimeError – 如果张量是需要梯度的叶子节点

__imul__(other)

张量原地乘法运算符，等价于 *=。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 另一个张量或标量值
返回:: 原地修改后的张量
返回类型:: riemann.TN
抛出:: RuntimeError – 如果张量是需要梯度的叶子节点

__itruediv__(other)

张量原地除法运算符，等价于 /=。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 另一个张量或标量值
返回:: 原地修改后的张量
返回类型:: riemann.TN
抛出:: RuntimeError – 如果张量是需要梯度的叶子节点

__ipow__(other)

张量原地幂运算符，等价于 **=。

参数:: other (riemann.TN 或 int 或 float 或 complex) – 指数张量或标量值
返回:: 原地修改后的张量
返回类型:: riemann.TN
抛出:: RuntimeError – 如果张量是需要梯度的叶子节点

数学运算

riemann.matmul(input, other)

两个张量的矩阵乘法。

参数:

input (riemann.TN) – 第一个张量
other (riemann.TN) – 第二个张量

返回:

张量的矩阵乘积

返回类型:

riemann.TN

riemann.dot(x, y)

计算两个张量的点积。

参数:

x (riemann.TN) – 第一个张量
y (riemann.TN) – 第二个张量

返回:

点积结果

返回类型:

riemann.TN

riemann.outer(x, y)

计算两个张量的外积。

参数:

x (riemann.TN) – 第一个张量
y (riemann.TN) – 第二个张量

返回:

外积结果

返回类型:

riemann.TN

riemann.cross(input, other, dim=-1)

计算两个张量的矢量积（叉积）。

参数:

input (riemann.TN) – 第一个张量
other (riemann.TN) – 第二个张量
dim (int, optional) – 进行叉积计算的维度，默认为最后一个维度

返回:

矢量积结果

返回类型:

riemann.TN

riemann.einsum(equation, *operands)

使用爱因斯坦求和约定计算张量运算。

参数:

equation (str) – 爱因斯坦求和方程字符串，如 “ij,jk->ik”
operands (riemann.TN) – 输入张量序列

返回:

计算结果张量

返回类型:

riemann.TN

抛出:

ValueError – 如果方程格式无效或张量维度不匹配

riemann.trace(input)

返回输入2-D矩阵对角线元素的和。

参数:: input (riemann.TN) – 输入的2-D张量
返回:: 对角线元素之和的标量张量
返回类型:: riemann.TN
抛出:: ValueError – 如果输入不是2-D张量

riemann.kron(input, other, *, out=None)

计算 input 和 other 的克罗内克积。

参数:

input (riemann.TN) – 第一个输入张量
other (riemann.TN) – 第二个输入张量
out (riemann.TN, optional) – 可选的输出张量，用于存储结果

返回:

克罗内克积结果张量

返回类型:

riemann.TN

抛出:

ValueError – 如果两个张量的维度不同

riemann.sum(x, dim=None, keepdim=False)

计算跨维度元素的总和。

参数:

x (riemann.TN) – 输入张量
dim (int 或整数元组, optional) – 要求和的维度
keepdim (bool, optional) – 是否保持缩减的维度

返回:

元素的总和

返回类型:

riemann.TN

riemann.prod(x, dim=None, keepdim=False)

计算跨维度元素的乘积。

参数:

x (riemann.TN) – 输入张量
dim (int 或整数元组, optional) – 要相乘的维度
keepdim (bool, optional) – 是否保持缩减的维度

返回:

元素的乘积

返回类型:

riemann.TN

riemann.mean(x, dim=None, keepdim=False)

计算跨维度元素的平均值。

参数:

x (riemann.TN) – 输入张量
dim (int 或整数元组, optional) – 要平均的维度
keepdim (bool, optional) – 是否保持缩减的维度

返回:

元素的平均值

返回类型:

riemann.TN

riemann.var(x, dim=None, unbiased=True, keepdim=False)

计算跨维度元素的方差。

参数:

x (riemann.TN) – 输入张量
dim (int 或整数元组, optional) – 要计算方差的维度
unbiased (bool, optional) – 是否使用无偏估计
keepdim (bool, optional) – 是否保持缩减的维度

返回:

元素的方差

返回类型:

riemann.TN

riemann.std(x, dim=None, unbiased=True, keepdim=False)

计算跨维度元素的标准差。

参数:

x (riemann.TN) – 输入张量
dim (int 或整数元组, optional) – 要计算标准差的维度
unbiased (bool, optional) – 是否使用无偏估计
keepdim (bool, optional) – 是否保持缩减的维度

返回:

元素的标准差

返回类型:

riemann.TN

riemann.norm(x, p='fro', dim=None, keepdim=False)

计算张量的范数。

参数:

x (riemann.TN) – 输入张量
p (int, float, str, optional) – 范数阶数
dim (int 或整数元组, optional) – 要计算范数的维度
keepdim (bool, optional) – 是否保持缩减的维度

返回:

张量的范数

返回类型:

riemann.TN

riemann.max(x, dim=None, keepdim=False, *, out=None)

计算跨维度元素的最大值。

参数:

x (riemann.TN) – 输入张量
dim (int, optional) – 要查找最大值的维度
keepdim (bool, optional) – 是否保持缩减的维度
out (riemann.TN, optional) – 输出张量

返回:

最大值

返回类型:

riemann.TN

riemann.min(x, dim=None, keepdim=False, *, out=None)

计算跨维度元素的最小值。

参数:

x (riemann.TN) – 输入张量
dim (int, optional) – 要查找最小值的维度
keepdim (bool, optional) – 是否保持缩减的维度
out (riemann.TN, optional) – 输出张量

返回:

最小值

返回类型:

riemann.TN

riemann.argmax(x, dim=None, keepdim=False, *, out=None)

计算跨维度元素的最大值索引。

参数:

x (riemann.TN) – 输入张量
dim (int, optional) – 要查找最大值索引的维度
keepdim (bool, optional) – 是否保持缩减的维度
out (riemann.TN, optional) – 输出张量

返回:

最大值索引

返回类型:

riemann.TN

riemann.argmin(x, dim=None, keepdim=False, *, out=None)

计算跨维度元素的最小值索引。

参数:

x (riemann.TN) – 输入张量
dim (int, optional) – 要查找最小值索引的维度
keepdim (bool, optional) – 是否保持缩减的维度
out (riemann.TN, optional) – 输出张量

返回:

最小值索引

返回类型:

riemann.TN

riemann.abs(x)

计算每个元素的绝对值。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的绝对值
返回类型:: riemann.TN

riemann.sqrt(x)

计算每个元素的平方根。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的平方根
返回类型:: riemann.TN

riemann.pow(input, exponent)

将每个元素提升到幂。

参数:

input (riemann.TN) – 输入张量
exponent (riemann.TN 或标量) – 指数值

返回:

输入张量的幂

返回类型:

riemann.TN

riemann.log(x)

计算每个元素的自然对数。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的自然对数
返回类型:: riemann.TN

riemann.log1p(x)

计算每个元素加一的自然对数。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素加一的自然对数
返回类型:: riemann.TN

riemann.log2(x)

计算每个元素的以2为底的对数。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的以2为底的对数
返回类型:: riemann.TN

riemann.log10(x)

计算每个元素的以10为底的对数。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的以10为底的对数
返回类型:: riemann.TN

riemann.exp(x)

计算每个元素的指数。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的指数
返回类型:: riemann.TN

riemann.exp2(x)

计算每个元素的2的幂。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的2的幂
返回类型:: riemann.TN

riemann.square(x)

计算每个元素的平方。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的平方
返回类型:: riemann.TN

riemann.sin(x)

计算每个元素的正弦。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的正弦
返回类型:: riemann.TN

riemann.cos(x)

计算每个元素的余弦。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的余弦
返回类型:: riemann.TN

riemann.tan(x)

计算每个元素的正切。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的正切
返回类型:: riemann.TN

riemann.cot(x)

计算每个元素的余切。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的余切
返回类型:: riemann.TN

riemann.sec(x)

计算每个元素的正割。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的正割
返回类型:: riemann.TN

riemann.csc(x)

计算每个元素的余割。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的余割
返回类型:: riemann.TN

riemann.asin(x)

计算每个元素的反正弦。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的反正弦
返回类型:: riemann.TN

riemann.acos(x)

计算每个元素的反余弦。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的反余弦
返回类型:: riemann.TN

riemann.atan(x)

计算每个元素的反正切。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的反正切
返回类型:: riemann.TN

riemann.sinh(x)

计算每个元素的双曲正弦。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的双曲正弦
返回类型:: riemann.TN

riemann.cosh(x)

计算每个元素的双曲余弦。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的双曲余弦
返回类型:: riemann.TN

riemann.tanh(x)

计算每个元素的双曲正切。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的双曲正切
返回类型:: riemann.TN

riemann.coth(x)

计算每个元素的双曲余切。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的双曲余切
返回类型:: riemann.TN

riemann.sech(x)

计算每个元素的双曲正割。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的双曲正割
返回类型:: riemann.TN

riemann.csch(x)

计算每个元素的双曲余割。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的双曲余割
返回类型:: riemann.TN

riemann.arcsinh(x)

计算每个元素的反双曲正弦。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的反双曲正弦
返回类型:: riemann.TN

riemann.arccosh(x)

计算每个元素的反双曲余弦。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的反双曲余弦
返回类型:: riemann.TN

riemann.arctanh(x)

计算每个元素的反双曲正切。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的反双曲正切
返回类型:: riemann.TN

riemann.ceil(x)

向上取整，返回不小于每个元素的最小整数。

参数:: x (riemann.TN) – 输入张量
返回:: 向上取整后的张量
返回类型:: riemann.TN

riemann.floor(x)

向下取整，返回不大于每个元素的最大整数。

参数:: x (riemann.TN) – 输入张量
返回:: 向下取整后的张量
返回类型:: riemann.TN

riemann.round(x)

四舍五入到最近的整数。

参数:: x (riemann.TN) – 输入张量
返回:: 四舍五入后的张量
返回类型:: riemann.TN

riemann.trunc(x)

截断小数部分，返回整数部分。

参数:: x (riemann.TN) – 输入张量
返回:: 截断后的张量
返回类型:: riemann.TN

riemann.sign(x)

计算每个元素的符号。

参数:: x (riemann.TN) – 输入张量
返回:: 每个元素的符号
返回类型:: riemann.TN

riemann.where(cond, x=None, y=None)

根据条件从 x 或 y 中选择元素。

参数:

cond (riemann.TN) – 条件张量
x (riemann.TN, optional) – 当条件为 True 时选择的张量
y (riemann.TN, optional) – 当条件为 False 时选择的张量

返回:

从 x 或 y 中选取元素组成的张量

返回类型:

riemann.TN

riemann.clamp(x, min=None, max=None, out=None)

将所有元素限制在指定范围内。

参数:

x (riemann.TN) – 输入张量
min (float, optional) – 最小值
max (float, optional) – 最大值
out (riemann.TN, optional) – 输出张量

返回:

元素被限制在指定范围内的张量

返回类型:

riemann.TN

riemann.masked_fill(input, mask, value)

根据掩码填充值到张量中。

参数:

input (riemann.TN) – 输入张量
mask (riemann.TN) – 掩码张量，形状与输入张量相同
value (scalar) – 填充的值

返回:

填充后的张量

返回类型:

riemann.TN

riemann.maximum(input, other)

计算两个张量的元素级最大值。

参数:

input (riemann.TN) – 第一个输入张量
other (riemann.TN) – 第二个输入张量

返回:

元素级最大值组成的张量

返回类型:

riemann.TN

riemann.minimum(input, other)

计算两个张量的元素级最小值。

参数:

input (riemann.TN) – 第一个输入张量
other (riemann.TN) – 第二个输入张量

返回:

元素级最小值组成的张量

返回类型:

riemann.TN

riemann.diagonal(input, offset=0, dim1=-2, dim2=-1)

返回张量的对角线。

参数:

input (riemann.TN) – 输入张量
offset (int, optional) – 对角线的偏移量
dim1 (int, optional) – 对角线的第一个维度
dim2 (int, optional) – 对角线的第二个维度

返回:

张量的对角线

返回类型:

riemann.TN

riemann.diag(input, offset=0)

返回二维张量的对角线或构造对角矩阵。

参数:

input (riemann.TN) – 输入张量
offset (int, optional) – 对角线的偏移量

返回:

张量的对角线或对角矩阵

返回类型:

riemann.TN

riemann.fill_diagonal(input, value, offset=0, dim1=-2, dim2=-1)

用指定值填充张量的对角线。

参数:

input (riemann.TN) – 输入张量
value (标量) – 填充对角线的值
offset (int, optional) – 对角线的偏移量
dim1 (int, optional) – 对角线的第一个维度
dim2 (int, optional) – 对角线的第二个维度

返回:

对角线被填充的张量

返回类型:

riemann.TN

riemann.batch_diag(v)

返回张量的批处理对角线。

参数:: v (riemann.TN) – 输入张量
返回:: 张量的批处理对角线
返回类型:: riemann.TN

riemann.nonzero(input, *, as_tuple=False)

返回非零元素的索引。

参数:

input (riemann.TN) – 输入张量
as_tuple (bool, optional) – 是否以张量元组的形式返回

返回:

非零元素的索引

返回类型:

riemann.TN 或 riemann.TN 的元组

riemann.tril(input_tensor, diagonal=0)

返回矩阵的下三角部分。

参数:

input_tensor (riemann.TN) – 输入张量
diagonal (int, optional) – 对角线偏移量

返回:

矩阵的下三角部分

返回类型:

riemann.TN

riemann.triu(input_tensor, diagonal=0)

返回矩阵的上三角部分。

参数:

input_tensor (riemann.TN) – 输入张量
diagonal (int, optional) – 对角线偏移量

返回:

矩阵的上三角部分

返回类型:

riemann.TN

riemann.cumsum(input, dim, *, dtype=None)

计算张量沿指定维度的累积和。

参数:

input (riemann.TN) – 输入张量
dim (int) – 计算累积和的维度
dtype (riemann.dtype, optional) – 输出张量的数据类型

返回:

累积和结果

返回类型:

riemann.TN

riemann.unique(input, sorted=True, return_inverse=False, return_counts=False, dim=None)

返回张量中的唯一元素。

参数:

input (riemann.TN) – 输入张量
sorted (bool, optional) – 是否对唯一值进行排序
return_inverse (bool, optional) – 是否返回逆索引
return_counts (bool, optional) – 是否返回每个唯一值的计数
dim (int, optional) – 沿哪个维度查找唯一值，默认为None（展平后查找）

返回:

唯一值，如果指定return_inverse或return_counts则返回元组

返回类型:

riemann.TN or tuple

riemann.broadcast_tensors(*tensors)

广播多个张量到相同的形状。

参数:: tensors (riemann.TN) – 要广播的张量序列
返回:: 广播后的张量列表
返回类型:: list of riemann.TN

riemann.repeat(input, repeats, dim=None)

沿指定维度重复张量元素。

参数:

input (riemann.TN) – 输入张量
repeats (int) – 每个元素的重复次数
dim (int, optional) – 要重复的维度，默认为None（展平后重复）

返回:

重复后的张量

返回类型:

riemann.TN

比较运算

riemann.equal(a, b)

计算元素级的相等性。

参数:

a (riemann.TN) – 第一个张量
b (riemann.TN) – 第二个张量

返回:

指示相等性的布尔张量

返回类型:

bool

riemann.not_equal(a, b)

计算元素级的不等性。

参数:

a (riemann.TN) – 第一个张量
b (riemann.TN) – 第二个张量

返回:

指示不等性的布尔张量

返回类型:

bool

riemann.allclose(a, b, rtol=1e-5, atol=1e-8, equal_nan=False)

如果两个张量在容差范围内元素级相等，则返回 True。

参数:

a (riemann.TN) – 第一个张量
b (riemann.TN) – 第二个张量
rtol (float, optional) – 相对容差
atol (float, optional) – 绝对容差
equal_nan (bool, optional) – 是否将 NaN 值视为相等

返回:

张量是否接近

返回类型:

bool

riemann.isinf(x)

元素级测试是否为无穷大。

参数:: x (riemann.TN) – 输入张量
返回:: 指示无穷大的布尔张量
返回类型:: riemann.TN

riemann.isnan(x)

元素级测试是否为 NaN。

参数:: x (riemann.TN) – 输入张量
返回:: 指示 NaN 的布尔张量
返回类型:: riemann.TN

riemann.isreal(x)

元素级测试是否为实数。

参数:: x (riemann.TN) – 输入张量
返回:: 指示实数的布尔张量
返回类型:: riemann.TN

排序运算

riemann.sort(input, dim=-1, descending=False, stable=False, *, out=None)

沿给定维度对张量元素进行排序。

参数:

input (riemann.TN) – 输入张量
dim (int, optional) – 排序的维度
descending (bool, optional) – 是否按降序排序
stable (bool, optional) – 是否使用稳定排序
out (riemann.TN, optional) – 输出张量

返回:

排序后的张量和索引

返回类型:

riemann.TN, riemann.TN

riemann.argsort(input, dim=-1, descending=False, stable=False, *, out=None)

沿给定维度返回将张量排序的索引。

参数:

input (riemann.TN) – 输入张量
dim (int, optional) – 排序的维度
descending (bool, optional) – 是否按降序排序
stable (bool, optional) – 是否使用稳定排序
out (riemann.TN, optional) – 输出张量

返回:

排序索引

返回类型:

riemann.TN

原地操作

riemann.TN.setat_(index, val)

原地设置张量指定位置的值。

参数:

index (int, slice, tuple, or array) – 索引，指定要设置值的位置
val (riemann.TN, numpy.ndarray, list, or scalar) – 要设置的值

返回:

原地修改后的张量

返回类型:

riemann.TN

riemann.TN.addat_(index, val)

原地将值加到张量指定位置。

参数:

index (int, slice, tuple, or array) – 索引，指定要操作的位置
val (riemann.TN, numpy.ndarray, list, or scalar) – 要添加的值

返回:

原地修改后的张量

返回类型:

riemann.TN

riemann.TN.subat_(index, val)

原地从张量指定位置减去值。

参数:

index (int, slice, tuple, or array) – 索引，指定要操作的位置
val (riemann.TN, numpy.ndarray, list, or scalar) – 要减去的值

返回:

原地修改后的张量

返回类型:

riemann.TN

riemann.TN.mulat_(index, val)

原地将张量指定位置的值乘以给定值。

参数:

index (int, slice, tuple, or array) – 索引，指定要操作的位置
val (riemann.TN, numpy.ndarray, list, or scalar) – 要乘的值

返回:

原地修改后的张量

返回类型:

riemann.TN

riemann.TN.divat_(index, val)

原地将张量指定位置的值除以给定值。

参数:

index (int, slice, tuple, or array) – 索引，指定要操作的位置
val (riemann.TN, numpy.ndarray, list, or scalar) – 除数

返回:

原地修改后的张量

返回类型:

riemann.TN

riemann.TN.powat_(index, val)

原地对张量指定位置的值进行幂运算。

参数:

index (int, slice, tuple, or array) – 索引，指定要操作的位置
val (riemann.TN, numpy.ndarray, list, or scalar) – 指数

返回:

原地修改后的张量

返回类型:

riemann.TN

riemann.TN.scatter_(dim, index, src=None, *, value=None)

原地将值按照索引填充到张量中。

参数:

dim (int) – 沿着哪个维度进行索引
index (riemann.TN) – 索引张量
src (riemann.TN, optional) – 源张量，提供要填充的值
value (int, float, complex, optional) – 标量值，提供要填充的值

返回:

原地修改后的张量

返回类型:

riemann.TN

riemann.TN.scatter_add_(dim, index, src)

原地将值按照索引累加到张量中。

参数:

dim (int) – 沿着哪个维度进行索引
index (riemann.TN) – 索引张量
src (riemann.TN) – 源张量，提供要累加的值

返回:

原地修改后的张量

返回类型:

riemann.TN

riemann.TN.requires_grad_(requires_grad=True)

原地设置张量是否需要计算梯度。

参数:: requires_grad (bool, optional) – 是否需要计算梯度
返回:: 原地修改后的张量
返回类型:: riemann.TN

riemann.TN.add_(other)

原地加法操作。

参数:: other (riemann.TN, numpy.ndarray, list, or scalar) – 要加到当前张量上的值
返回:: 原地修改后的张量
返回类型:: riemann.TN

riemann.TN.sub_(other)

原地减法操作。

参数:: other (riemann.TN, numpy.ndarray, list, or scalar) – 要从当前张量中减去的值
返回:: 原地修改后的张量
返回类型:: riemann.TN

riemann.TN.mul_(other)

原地乘法操作。

参数:: other (riemann.TN, numpy.ndarray, list, or scalar) – 要与当前张量相乘的值
返回:: 原地修改后的张量
返回类型:: riemann.TN

riemann.TN.div_(other)

原地除法操作。

参数:: other (riemann.TN, numpy.ndarray, list, or scalar) – 除数
返回:: 原地修改后的张量
返回类型:: riemann.TN

riemann.TN.pow_(other)

原地幂运算操作。

参数:: other (riemann.TN, numpy.ndarray, list, or scalar) – 指数
返回:: 原地修改后的张量
返回类型:: riemann.TN

riemann.TN.detach_()

原地断开张量与计算图的连接。

返回:: 原地修改后的张量
返回类型:: riemann.TN

riemann.TN.copy_(src)

原地复制源张量到当前张量。

参数:: src (riemann.TN, numpy.ndarray, list, or scalar) – 源张量
返回:: 原地修改后的张量
返回类型:: riemann.TN

riemann.TN.zero_()

原地将张量所有元素设置为0。

返回:: 原地修改后的张量
返回类型:: riemann.TN

riemann.TN.fill_(value)

原地将张量所有元素填充为指定值。

参数:: value (riemann.TN, numpy.ndarray, list, or scalar) – 填充值
返回:: 原地修改后的张量
返回类型:: riemann.TN

riemann.TN.clamp_(min=None, max=None)

原地将张量元素限制在指定范围内。

参数:

min (float, optional) – 最小值
max (float, optional) – 最大值

返回:

原地修改后的张量

返回类型:

riemann.TN

riemann.masked_fill_(input, mask, value)

原地版本的 masked_fill 函数，根据掩码填充值到张量中。

参数:

input (riemann.TN) – 输入张量
mask (riemann.TN) – 掩码张量，形状与输入张量相同
value (scalar) – 填充的值

返回:

原地修改后的张量

返回类型:

riemann.TN

riemann.fill_diagonal_(input, value, offset=0, dim1=-2, dim2=-1)

fill_diagonal 的原地版本。

参数:

input (riemann.TN) – 输入张量
value (标量) – 填充对角线的值
offset (int, optional) – 对角线的偏移量
dim1 (int, optional) – 对角线的第一个维度
dim2 (int, optional) – 对角线的第二个维度

返回:

对角线被填充的输入张量

返回类型:

riemann.TN

收集与散射函数

riemann.TN.gather(dim, index)

根据指定的维度和索引收集元素。

参数:

dim (int) – 收集维度
index (riemann.TN) – 索引张量

返回:

收集后的张量

返回类型:

riemann.TN

riemann.TN.scatter(dim, index, src=None, *, value=None)

将值按照索引填充到新张量中。

参数:

dim (int) – 沿着哪个维度进行索引
index (riemann.TN) – 索引张量
src (riemann.TN, optional) – 源张量，提供要填充的值
value (int, float, complex, optional) – 标量值，提供要填充的值

返回:

填充后的新张量

返回类型:

riemann.TN

riemann.TN.scatter_(dim, index, src=None, *, value=None)

原地将值按照索引填充到张量中。

参数:

dim (int) – 沿着哪个维度进行索引
index (riemann.TN) – 索引张量
src (riemann.TN, optional) – 源张量，提供要填充的值
value (int, float, complex, optional) – 标量值，提供要填充的值

返回:

原地修改后的张量

返回类型:

riemann.TN

riemann.TN.scatter_add(dim, index, src)

将值按照索引累加到新张量中。

参数:

dim (int) – 沿着哪个维度进行索引
index (riemann.TN) – 索引张量
src (riemann.TN) – 源张量，提供要累加的值

返回:

累加后的新张量

返回类型:

riemann.TN

riemann.TN.scatter_add_(dim, index, src)

原地将值按照索引累加到张量中。

参数:

dim (int) – 沿着哪个维度进行索引
index (riemann.TN) – 索引张量
src (riemann.TN) – 源张量，提供要累加的值

返回:

原地修改后的张量

返回类型:

riemann.TN

数据转换

riemann.from_numpy(arr, requires_grad=False)

将 NumPy 数组转换为 Riemann 张量。

参数:

arr (numpy.ndarray) – 输入 NumPy 数组
requires_grad (bool, optional) – 是否跟踪此张量上的操作

返回:

Riemann 张量

返回类型:

riemann.TN

riemann.item(tensor)

将单个元素的张量转换为 Python 标量。

参数:: tensor (riemann.TN) – 输入张量
返回:: Python 标量
返回类型:: int, float, etc.

riemann.TN.tolist()

将张量转换为 Python 列表。

返回:: Python 列表或标量
返回类型:: list, int, float, complex

riemann.TN.numpy()

将张量转换为 NumPy 数组。

返回:: NumPy 数组
返回类型:: numpy.ndarray

riemann.TN.to(*args, **kwargs)

将张量转换为指定的数据类型和/或设备。

参数:

dtype (dtype, optional) – 目标数据类型
device (str, Device, optional) – 目标设备

返回:

转换后的张量

返回类型:

riemann.TN

riemann.TN.type(dtype=None)

返回或转换张量的数据类型。

参数:: dtype (dtype, optional) – 数据类型，如果为None则返回当前数据类型
返回:: 如果dtype为None返回当前数据类型，否则返回转换后的数据类型
返回类型:: dtype or riemann.TN

riemann.TN.type_as(other)

将张量转换为与另一个张量相同的数据类型。

参数:: other (riemann.TN) – 目标数据类型的参考张量
返回:: 转换后的数据类型
返回类型:: riemann.TN

riemann.TN.bool()

将张量转换为布尔类型。

返回:: 布尔类型的张量
返回类型:: riemann.TN

riemann.TN.float()

将张量转换为单精度浮点类型(float32)。

返回:: float32类型的张量
返回类型:: riemann.TN

riemann.TN.double()

将张量转换为双精度浮点类型(float64)。

返回:: float64类型的张量
返回类型:: riemann.TN

副本函数

riemann.clone(tensor)

返回张量的副本。

参数:: tensor (riemann.TN) – 输入张量
返回:: 张量副本
返回类型:: riemann.TN

riemann.TN.copy()

返回张量的副本，不共享内存，也不依赖原张量。

返回:: 张量副本
返回类型:: riemann.TN

riemann.detach(tensor)

从计算图中分离张量，停止梯度跟踪。

参数:: tensor (riemann.TN) – 输入张量
返回:: 分离后的张量
返回类型:: riemann.TN

数据类型

预定义数据类型

riemann.dtype.float16: 16位浮点数据类型

riemann.dtype.float32: 32位浮点数据类型

riemann.dtype.float64: 64位浮点数据类型

riemann.dtype.complex64: 64位复数数据类型

riemann.dtype.complex128: 128位复数数据类型

riemann.dtype.half: float16的别名

riemann.dtype.float_: float32的别名

riemann.dtype.double: float64的别名

riemann.dtype.complex_: complex64的别名

riemann.dtype.int8: 8位有符号整数数据类型

riemann.dtype.int16: 16位有符号整数数据类型

riemann.dtype.int32: 32位有符号整数数据类型

riemann.dtype.int64: 64位有符号整数数据类型

riemann.dtype.uint8: 8位无符号整数数据类型

riemann.dtype.uint16: 16位无符号整数数据类型

riemann.dtype.uint32: 32位无符号整数数据类型

riemann.dtype.uint64: 64位无符号整数数据类型

riemann.dtype.short: int16的别名

riemann.dtype.int_: int32的别名

riemann.dtype.long: int64的别名

riemann.dtype.bool_: 布尔数据类型

默认数据类型函数

riemann.set_default_dtype(dtype)

设置默认的浮点数据类型

参数:: dtype (numpy.dtype) – 要设置为默认的浮点数据类型

riemann.get_default_dtype()

获取当前默认的浮点数据类型

返回:: 当前默认的浮点数据类型
返回类型:: numpy.dtype

riemann.get_default_complex()

根据默认的浮点类型推导默认的复数数据类型

返回:: 默认的复数数据类型
返回类型:: numpy.dtype

数据类型检查函数

riemann.is_numeric_array(numpy_arr)

检查NumPy数组是否具有数值数据类型

参数:: numpy_arr (numpy.ndarray) – 要检查的NumPy数组
返回:: 数组是否具有数值数据类型
返回类型:: bool

riemann.is_number(v)

检查值是否为数值类型

参数:: v (Any) – 要检查的值
返回:: 值是否为数值类型
返回类型:: bool

riemann.is_float_or_complex(dtype)

检查数据类型是否为浮点或复数类型

参数:: dtype (numpy.dtype) – 要检查的数据类型
返回:: 数据类型是否为浮点或复数类型
返回类型:: bool

riemann.is_scalar(value)

检查值是否为标量（包括Riemann张量标量）

参数:: value (Any) – 要检查的值
返回:: 值是否为标量
返回类型:: bool

数据类型推断

riemann.infer_data_type(v)

从Python值、NumPy数组或值集合中推断适当的数据类型

参数:: v (Any) – 要推断数据类型的值或值集合
返回:: 推断的数据类型
返回类型:: numpy.dtype

梯度模式控制

riemann.is_grad_enabled()

获取当前线程的梯度计算状态

返回:: 当前梯度计算模式（True表示启用，False表示禁用）
返回类型:: bool

riemann.no_grad(func=None)

上下文管理器，用于暂时禁用梯度计算

也可以作为函数装饰器使用，禁用被装饰函数内所有计算的梯度追踪。

参数:: func (callable, optional) – 可选，如果提供，则将no_grad作为装饰器应用于该函数
返回:: 如果未提供func，则返回上下文管理器实例；如果提供了func，则返回装饰后的函数

示例：

# 作为上下文管理器使用
with riemann.no_grad():
    # 这段代码中的计算不会追踪梯度
    output = model(input_data)

# 作为装饰器使用
@riemann.no_grad
def inference(x):
    # 函数内的计算不会追踪梯度
    return model(x)

riemann.enable_grad(func=None)

上下文管理器，用于暂时启用梯度计算

也可以作为函数装饰器使用，确保被装饰函数内的计算追踪梯度。

参数:: func (callable, optional) – 可选，如果提供，则将enable_grad作为装饰器应用于该函数
返回:: 如果未提供func，则返回上下文管理器实例；如果提供了func，则返回装饰后的函数

示例：

# 作为上下文管理器使用
with riemann.enable_grad():
    # 这段代码中的计算会追踪梯度
    output = model(input_data)
    loss = loss_fn(output, target)
    loss.backward()

# 作为装饰器使用
@riemann.enable_grad
def train_step(x, y):
    # 函数内的计算会追踪梯度
    pred = model(x)
    loss = loss_fn(pred, y)
    loss.backward()
    return loss

riemann.set_grad_enabled(mode=True, func=None)

上下文管理器，用于设置梯度计算模式

类似于PyTorch的set_grad_enabled()，可以显式地启用或禁用梯度计算。支持作为上下文管理器或装饰器使用，提供最灵活的梯度控制方式。

参数:

mode (bool) – 如果为True，则启用梯度计算；如果为False，则禁用梯度计算
func (callable, optional) – 可选，当作为装饰器使用时传入的函数

返回:

如果func为None，返回上下文管理器实例；如果提供了func参数，返回包装后的函数

示例：

# 作为上下文管理器使用
with riemann.set_grad_enabled(False):
    # 这段代码中的计算不会追踪梯度
    output = model(input_data)

with riemann.set_grad_enabled(True):
    # 这段代码中的计算会追踪梯度
    output = model(input_data)
    loss = loss_fn(output, target)
    loss.backward()

# 作为装饰器使用
@riemann.set_grad_enabled(False)
def inference(x):
    return model(x)

@riemann.set_grad_enabled(True)
def train(x, y):
    pred = model(x)
    loss = loss_fn(pred, y)
    loss.backward()
    return loss

序列化

riemann.save(obj, f, pickle_module=None, pickle_protocol=2, use_new_zipfile_serialization=True)

将对象保存到磁盘文件。

此函数使用pickle序列化将Riemann张量、参数、模块或任何Python对象保存到磁盘文件。

参数:

obj (Any) – 要保存的对象。可以是张量、参数、模块或任何可pickle的对象
f (str, os.PathLike, 或类文件对象) – 要写入的文件路径或类文件对象
pickle_module (Any, optional) – 用于pickle的模块（默认：pickle）
pickle_protocol (int, optional) – Pickle协议版本（默认：2）
use_new_zipfile_serialization (bool, optional) – 是否使用基于zip的序列化（默认：True）

示例：

>>> import riemann as rm
>>> # 保存张量
>>> tensor = rm.randn(3, 4)
>>> rm.save(tensor, 'tensor.pt')
>>>
>>> # 保存模块
>>> model = rm.nn.Linear(10, 5)
>>> rm.save(model.state_dict(), 'model_weights.pt')
>>>
>>> # 保存多个对象
>>> rm.save({
...     'model': model.state_dict(),
...     'optimizer_state': optimizer.state_dict(),
...     'epoch': 10
... }, 'checkpoint.pt')

riemann.load(f, map_location=None, pickle_module=None, **pickle_load_args)

从磁盘文件加载对象。

此函数使用pickle反序列化从磁盘文件加载Riemann张量、参数、模块或任何Python对象。

参数:

f (str, os.PathLike, 或类文件对象) – 要读取的文件路径或类文件对象
map_location (Any, optional) – 用于重新映射存储位置的函数或字典
pickle_module (Any, optional) – 用于unpickle的模块（默认：pickle）
**pickle_load_args – 传递给pickle.load的额外参数

返回:

加载的对象

示例：

>>> import riemann as rm
>>> # 加载张量
>>> tensor = rm.load('tensor.pt')
>>>
>>> # 加载模型权重
>>> state_dict = rm.load('model_weights.pt')
>>> model.load_state_dict(state_dict)
>>>
>>> # 加载检查点
>>> checkpoint = rm.load('checkpoint.pt')
>>> model.load_state_dict(checkpoint['model'])
>>> optimizer.load_state_dict(checkpoint['optimizer_state'])
>>> epoch = checkpoint['epoch']

CUDA支持

class riemann.cuda.Device(device='cpu')

表示一个设备（CPU或CUDA GPU）。

参数:: device (str 或 int) – 设备类型或索引。可以是： - 字符串：’cpu’、’cuda’或’cuda:0’、’cuda:1’ - 整数：CUDA设备索引

示例：

>>> import riemann as rm
>>> # 创建CPU设备
>>> cpu_device = rm.Device('cpu')
>>> # 创建CUDA设备
>>> cuda_device = rm.Device('cuda')
>>> # 创建特定的CUDA设备
>>> cuda_device_1 = rm.Device('cuda:1')
>>> # 通过type和index参数创建CUDA设备
>>> cuda_device_2 = rm.Device('cuda', 2)

__enter__(): 进入设备上下文。

__exit__(exc_type, exc_val, exc_tb): 退出设备上下文。

__eq__(other): 与另一个设备比较。

__str__(): 返回设备的字符串表示。

__repr__(): 返回设备的官方字符串表示。

riemann.cuda.is_available()

检查CUDA是否可用。

返回:: 如果CUDA可用返回True，否则返回False
返回类型:: bool

riemann.cuda.device_count()

返回可用的CUDA设备数量。

返回:: 可用的CUDA设备数量
返回类型:: int

riemann.cuda.current_device()

返回当前CUDA设备的索引。

返回:: 当前CUDA设备的索引
返回类型:: int

riemann.cuda.get_device_name(device_idx)

返回给定索引的CUDA设备名称。

参数:: device_idx (int) – CUDA设备的索引
返回:: CUDA设备的名称
返回类型:: str

riemann.cuda.set_device(device_idx)

设置当前CUDA设备。

参数:: device_idx (int) – 要设置为当前的CUDA设备索引

riemann.cuda.empty_cache(): 清空CUDA缓存。

riemann.cuda.synchronize(device=None)

等待当前CUDA设备上的所有操作完成。

参数:: device (str, int, or Device, optional) – 要同步的设备，默认为当前设备

riemann.cuda.is_in_cuda_context()

检查当前线程是否在CUDA设备上下文中。

返回:: 如果在CUDA设备上下文中返回True，否则返回False
返回类型:: bool

riemann.memory_allocated(device_idx=None)

返回给定CUDA设备上分配的内存量。

参数:: device_idx (int, optional) – CUDA设备的索引。如果为None，使用当前设备
返回:: 分配的内存量（字节）
返回类型:: int

riemann.get_default_device()

获取张量创建的默认设备。

返回:: 默认设备
返回类型:: Device

riemann.set_default_device(device)

设置张量创建的默认设备。

参数:: device (str, int, 或 Device) – 要设置为默认的设备。可以是： - 字符串：’cpu’、’cuda’或’cuda:0’、’cuda:1’ - 整数：CUDA设备索引 - Device对象

示例：

>>> import riemann as rm
>>> rm.get_default_device()
device(type='cpu', index=None)
>>> rm.set_default_device('cuda')
>>> rm.get_default_device()
device(type='cuda', index=0)
>>> rm.set_default_device('cuda:1')
>>> rm.get_default_device()
device(type='cuda', index=1)

自动微分

梯度计算

riemann.autograd.backward(self, gradient=None, retain_graph=False, create_graph=False)

执行反向模式自动微分（反向传播）。

从当前张量开始，通过计算图向后传播梯度，为所有叶节点或设置了retains_grad=True的中间节点计算并存储梯度。

参数:

self (riemann.TN) – 触发反向传播的张量
gradient (riemann.TN or None, optional) – 输出张量的梯度，默认为None
retain_graph (bool, optional) – 此参数用于PyTorch兼容性，Riemann反向传播不依赖于此参数
create_graph (bool, optional) – 是否在梯度计算过程中创建计算图，设置为True以进行高阶导数计算

riemann.autograd.grad(outputs, inputs, grad_outputs=None, retain_graph=False, create_graph=False, allow_unused=False)

计算并返回输出相对于输入的梯度。

这是Riemann框架中的核心梯度计算函数。与backward()方法类似，但它直接返回计算出的梯度张量，而不是将它们存储在输入张量的.grad属性中。这使其更适合于高级梯度计算场景，如计算雅可比矩阵、海森矩阵等。

参数:

outputs (riemann.TN) – 要计算梯度的输出张量
inputs (riemann.TN or list/tuple of riemann.TN) – 要计算梯度的输入张量或输入张量列表/元组
grad_outputs (riemann.TN or None, optional) – 输出张量的梯度，默认为None
retain_graph (bool, optional) – 此参数用于PyTorch兼容性，Riemann反向传播不依赖于此参数
create_graph (bool, optional) – 是否在梯度计算过程中创建计算图
allow_unused (bool, optional) – 是否允许未使用的输入

返回:

对应于输入的梯度张量元组

返回类型:

tuple of riemann.TN

riemann.autograd.higher_order_grad(outputs, inputs, n, create_graph=False)

计算标量张量输出相对于inputs中每个张量的n阶导数。

此函数通过递归调用grad()来计算高阶导数。对于每个输入张量，它计算n阶导数并返回与输入列表对应的导数元组。

参数:

outputs (riemann.TN) – 要计算高阶导数的标量张量输出
inputs (riemann.TN or list/tuple of riemann.TN) – 要计算高阶导数的输入张量或输入张量列表/元组
n (int) – 导数的阶数，必须是非负整数
create_graph (bool, optional) – 是否在梯度计算过程中创建计算图

返回:

对应于输入的n阶导数张量元组

返回类型:

tuple of riemann.TN

riemann.autograd.gradcheck(func, inputs, eps=1e-6, atol=1e-5, rtol=1e-3, raise_exception=True, check_sparse_nnz=False, fast_mode=False)

通过比较数值梯度和解析梯度来验证给定函数的梯度计算是否正确。

此函数通过向输入参数添加小扰动来计算数值梯度（使用有限差分），并将其与使用自动微分计算的解析梯度进行比较。

参数:

func (callable) – 要验证梯度的函数
inputs (tuple of riemann.TN) – 测试用的输入张量元组
eps (float, optional) – 数值梯度计算的小扰动
atol (float, optional) – 绝对误差容差
rtol (float, optional) – 相对误差容差
raise_exception (bool, optional) – 如果梯度检查失败是否抛出异常
check_sparse_nnz (bool, optional) – 是否检查稀疏张量非零元素（当前版本不支持）
fast_mode (bool, optional) – 是否使用快速模式（仅检查第一个元素）

返回:

如果梯度检查通过则为True，否则为False

返回类型:

bool

riemann.track_grad(grad_func)

创建一个梯度跟踪装饰器，用于为函数添加自动微分支持。

这个装饰器工厂接收一个梯度函数，返回一个修饰器，该修饰器可以将普通的张量运算函数转换为支持自动微分的函数。它会自动创建反向传播函数，并管理梯度计算图的构建。

参数:: grad_func (callable) – 梯度计算函数，接收与前向函数相同的输入参数，返回一个元组，包含每个输入张量对应的梯度（偏导数）元组内元素需要与前向函数的输入张量一一对应，对于不需要梯度的张量，对应的梯度值应为None
返回:: 一个修饰器函数，用于包装前向计算函数，使其支持自动微分
返回类型:: callable

示例：

# 定义单输入导数函数（d/dx log(x) = 1/x）
def _log_derivative(x):
    return (1. / x.conj(),)

# 使用track_grad修饰器创建支持自动微分的对数函数
@track_grad(_log_derivative)
def mylog(x):
    return tensor(np.log(x.data))

# 使用带自动微分的对数函数
x = tensor(2., requires_grad=True)
y = mylog(x)
y.backward()
print(f'x.grad = {x.grad}')  # 输出: x.grad = 0.5

# 定义多输入导数函数（d/dx (x + y) = 1, d/dy (x + y) = 1）
def _add_derivative(x, y):
    return (tensor(1.), tensor(1.))

# 使用track_grad修饰器创建支持自动微分的加法函数
@track_grad(_add_derivative)
def myadd(x, y):
    return tensor(x.data + y.data)

# 使用带自动微分的加法函数
x = tensor(2., requires_grad=True)
y = tensor(3., requires_grad=True)
z = myadd(x, y)
z.backward()
print(f'x.grad = {x.grad}')  # 输出: x.grad = 1.0
print(f'y.grad = {y.grad}')  # 输出: y.grad = 1.0

class riemann.autograd.Function

Riemann框架中用于自定义梯度实现的基类，设计了与PyTorch的torch.autograd.Function类似的接口。

要使用此类，继承它并实现forward和backward静态方法： - forward: 执行前向计算，返回输出张量 - backward: 接收输出梯度，返回输入梯度

示例：

class MyFunction(Function):
    @staticmethod
    def forward(ctx, input1, input2):
        ctx.save_for_backward(input1, input2)
        output = input1 * input2
        return output

    @staticmethod
    def backward(ctx, grad_output):
        input1, input2 = ctx.saved_tensors
        grad_input1 = grad_output * input2
        grad_input2 = grad_output * input1
        return grad_input1, grad_input2

求导功能函数

riemann.autograd.functional.jacobian(func, inputs, create_graph=False, strict=True)

计算函数的雅可比矩阵。

该函数计算给定函数在输入点处的雅可比矩阵，支持单个或多个输入、单个或多个输出的情况，并与PyTorch的jacobian函数行为保持兼容。

参数:

func (callable) – 要计算雅可比矩阵的函数
inputs (riemann.TN or list/tuple of riemann.TN) – 函数的输入张量或张量列表/元组
create_graph (bool, optional) – 是否在梯度计算过程中创建计算图
strict (bool, optional) – 是否严格遵循PyTorch的行为规范

返回:

对应于输入/输出类型的雅可比矩阵表示

返回类型:

riemann.TN or list/tuple of riemann.TN

riemann.autograd.functional.hessian(func, inputs, create_graph=False, strict=True)

计算函数的 Hessian 矩阵。

该函数计算给定函数在输入点处的 Hessian 矩阵，即梯度的雅可比矩阵。它支持单个或多个输入的情况，并与PyTorch的hessian函数行为保持兼容。

参数:

func (callable) – 要计算 Hessian 矩阵的标量值函数
inputs (riemann.TN or list/tuple of riemann.TN) – 函数的输入张量或张量列表/元组
create_graph (bool, optional) – 是否在梯度计算过程中创建计算图
strict (bool, optional) – 如果为True，当检测到输出与输入无关时会引发错误

返回:

对应于输入类型的 Hessian 矩阵表示

返回类型:

riemann.TN or list/tuple of riemann.TN

riemann.autograd.functional.jvp(func, inputs, v=None, create_graph=False, strict=False)

计算 Jacobian 向量乘积（Jacobian-Vector Product）。

参数:

func (callable) – 要计算 JVP 的函数
inputs (riemann.TN or list/tuple of riemann.TN) – 函数的输入张量或张量列表/元组
v (riemann.TN or list/tuple of riemann.TN, optional) – 与 Jacobian 矩阵相乘的向量
create_graph (bool, optional) – 是否在梯度计算过程中创建计算图，用于高阶导数计算
strict (bool, optional) – 是否对未使用的输入引发错误

返回:

函数输出和 JVP 值

返回类型:

tuple of (riemann.TN, riemann.TN or list/tuple of riemann.TN)

riemann.autograd.functional.vjp(func, inputs, v=None, create_graph=False, strict=False)

计算向量 Jacobian 乘积（Vector-Jacobian Product）。

参数:

func (callable) – 要计算 VJP 的函数
inputs (riemann.TN or list/tuple of riemann.TN) – 函数的输入张量或张量列表/元组
v (riemann.TN or list/tuple of riemann.TN, optional) – 与 Jacobian 矩阵相乘的向量
create_graph (bool, optional) – 是否在梯度计算过程中创建计算图，用于高阶导数计算
strict (bool, optional) – 是否对未使用的输入引发错误

返回:

函数输出和 VJP 值

返回类型:

tuple of (riemann.TN, riemann.TN or list/tuple of riemann.TN)

riemann.autograd.functional.hvp(func, inputs, v, create_graph=False, strict=False)

计算 Hessian 向量乘积（Hessian-Vector Product）。

参数:

func (callable) – 要计算 HVP 的标量值函数
inputs (riemann.TN or list/tuple of riemann.TN) – 函数的输入张量或张量列表/元组
v (riemann.TN or list/tuple of riemann.TN) – 与 Hessian 矩阵相乘的向量
create_graph (bool, optional) – 是否在梯度计算过程中创建计算图，用于高阶导数计算
strict (bool, optional) – 是否对未使用的输入引发错误

返回:

函数输出和 HVP 值

返回类型:

tuple of (riemann.TN, riemann.TN or list/tuple of riemann.TN)

riemann.autograd.functional.vhp(func, inputs, v, create_graph=False, strict=False)

计算向量 Hessian 乘积（Vector-Hessian Product）。

参数:

func (callable) – 要计算 VHP 的标量值函数
inputs (riemann.TN or list/tuple of riemann.TN) – 函数的输入张量或张量列表/元组
v (riemann.TN or list/tuple of riemann.TN) – 与 Hessian 矩阵相乘的向量
create_graph (bool, optional) – 是否在梯度计算过程中创建计算图，用于高阶导数计算
strict (bool, optional) – 是否对未使用的输入引发错误

返回:

函数输出和 VHP 值

返回类型:

tuple of (riemann.TN, riemann.TN or list/tuple of riemann.TN)

riemann.autograd.functional.derivative(func, create_graph=False)

计算函数的导数函数。

该函数返回一个新函数，该新函数在调用时会计算原始函数func在输入点处的导数。支持func的输入为单个或多个张量，返回为单个或多个张量或标量。内部基于jacobian函数实现导数计算。

参数:

func (callable) – 要求导的函数
create_graph (bool, optional) – 是否在梯度计算中创建计算图，默认为False

返回:

导函数，该函数接受与原函数相同的输入

返回类型:

callable

上下文管理器

riemann.no_grad(): 上下文管理器，用于禁用梯度计算。在这个上下文中的操作不会被记录在计算图中。

riemann.enable_grad(): 上下文管理器，用于启用梯度计算。

riemann.set_grad_enabled(mode)

上下文管理器，根据 mode 参数启用或禁用梯度计算。

参数:: mode (bool) – 启用梯度计算为 True，禁用为 False

线性代数模块

riemann.linalg 模块提供了各种线性代数运算，包括矩阵乘法、分解、求解等。

矩阵运算

riemann.linalg.matmul(a, b)

计算两个张量的矩阵乘积。

参数:

a (riemann.TN) – 第一个输入张量
b (riemann.TN) – 第二个输入张量

返回:

矩阵乘积结果

返回类型:

riemann.TN

riemann.linalg.cross(a, b, dim=-1)

计算两个张量的叉积（向量积）。

参数:

a (riemann.TN) – 第一个输入张量
b (riemann.TN) – 第二个输入张量
dim (int, optional) – 计算叉积的维度，默认为-1

返回:

叉积结果

返回类型:

riemann.TN

范数计算

riemann.linalg.norm(A, ord=None, dim=None, keepdim=False)

计算张量或矩阵的范数。

参数:

A (riemann.TN) – 输入张量
ord (int or float or str, optional) – 范数的阶，默认为Frobenius范数
dim (int or tuple, optional) – 计算范数的维度，默认为None（计算所有元素的范数）
keepdim (bool, optional) – 是否保留输出的维度

返回:

范数值

返回类型:

riemann.TN

riemann.linalg.vector_norm(x, ord=2, dim=None, keepdim=False)

计算向量范数。

参数:

x (riemann.TN) – 输入张量
ord (float, optional) – 范数的阶，默认为2（L2范数）
dim (int or tuple, optional) – 计算范数的维度
keepdim (bool, optional) – 是否保留输出的维度

返回:

范数值

返回类型:

riemann.TN

riemann.linalg.matrix_norm(A, ord='fro', dim=(-2, -1), keepdim=False)

计算矩阵范数。

参数:

A (riemann.TN) – 输入张量
ord (str or int, optional) – 范数的阶，默认为’fro’（Frobenius范数）
dim (tuple, optional) – 矩阵所在的维度，默认为(-2, -1)
keepdim (bool, optional) – 是否保留输出的维度

返回:

矩阵范数值

返回类型:

riemann.TN

riemann.linalg.cond(A, p=None)

计算矩阵的条件数。

参数:

A (riemann.TN) – 输入矩阵
p (int or float or str, optional) – 范数类型，默认为None（2-范数条件数）

返回:

条件数

返回类型:

riemann.TN

riemann.linalg.svdvals(A)

计算矩阵的奇异值。

参数:: A (riemann.TN) – 输入矩阵
返回:: 奇异值
返回类型:: riemann.TN

矩阵分解

riemann.linalg.det(A)

计算矩阵的行列式。

参数:: A (riemann.TN) – 输入矩阵
返回:: 行列式值
返回类型:: riemann.TN

riemann.linalg.inv(A)

计算方阵的逆矩阵。

参数:: A (riemann.TN) – 输入方阵
返回:: 逆矩阵
返回类型:: riemann.TN

riemann.linalg.skew(A)

计算矩阵的斜对称部分。

参数:: A (riemann.TN) – 输入矩阵
返回:: 斜对称矩阵
返回类型:: riemann.TN

riemann.linalg.svd(A, full_matrices=True)

计算矩阵的奇异值分解（SVD）。

参数:

A (riemann.TN) – 输入矩阵
full_matrices (bool, optional) – 是否返回完整的U和Vh矩阵

返回:

(U, S, Vh)元组

返回类型:

tuple

riemann.linalg.pinv(A, rcond=1e-15)

计算矩阵的Moore-Penrose伪逆。

参数:

A (riemann.TN) – 输入矩阵
rcond (float, optional) – 奇异值阈值

返回:

伪逆矩阵

返回类型:

riemann.TN

特征值分解

riemann.linalg.eig(A)

计算方阵的特征值和特征向量。

参数:: A (riemann.TN) – 输入方阵
返回:: (特征值, 特征向量)元组
返回类型:: tuple

riemann.linalg.eigh(A, UPLO='L')

计算厄米特矩阵（或实对称矩阵）的特征值和特征向量。

参数:

A (riemann.TN) – 输入厄米特矩阵
UPLO (str, optional) – 指定使用上三角(‘U’)还是下三角(‘L’)部分

返回:

(特征值, 特征向量)元组

返回类型:

tuple

线性方程组求解

riemann.linalg.lstsq(A, b, rcond=None)

计算最小二乘解。

参数:

A (riemann.TN) – 系数矩阵
b (riemann.TN) – 右侧向量或矩阵
rcond (float, optional) – 奇异值阈值

返回:

最小二乘解

返回类型:

riemann.TN

riemann.linalg.lu(A, pivot=True)

计算矩阵的LU分解。

参数:

A (riemann.TN) – 输入矩阵
pivot (bool, optional) – 是否进行主元选取

返回:

(P, L, U)元组

返回类型:

tuple

riemann.linalg.solve(A, b)

求解线性方程组 Ax = b。

参数:

A (riemann.TN) – 系数矩阵
b (riemann.TN) – 右侧向量或矩阵

返回:

解向量或矩阵

返回类型:

riemann.TN

riemann.linalg.qr(A, mode='reduced')

计算矩阵的QR分解。

参数:

A (riemann.TN) – 输入矩阵
mode (str, optional) – 分解模式，’reduced’或’complete’

返回:

(Q, R)元组

返回类型:

tuple

riemann.linalg.cholesky(A, upper=False)

计算正定矩阵的Cholesky分解。

参数:

A (riemann.TN) – 输入正定矩阵
upper (bool, optional) – 是否返回上三角矩阵，默认为False（下三角）

返回:

Cholesky因子

返回类型:

riemann.TN

神经网络模块

基础类

class riemann.nn.Module

神经网络模块的基础类。

__init__(): 初始化模块。

forward(*args, **kwargs)

定义前向传播。

参数:

args – 输入参数
kwargs – 关键字参数

返回:

前向传播的输出

parameters()

获取所有可训练参数。

返回:: 参数列表
返回类型:: list of riemann.TN

class riemann.nn.Parameter(data=None, requires_grad=True)

可训练参数类，用于存储模型参数。

参数:

data (array_like, optional) – 参数数据
requires_grad (bool, optional) – 是否跟踪梯度

容器模块

class riemann.nn.Sequential(*modules)

按顺序应用多个模块的容器。

参数:: modules (list of riemann.Module) – 模块列表

class riemann.nn.ModuleList(modules=None)

用于存储模块列表的容器类。

该容器允许以列表形式存储多个模块，并提供便捷的访问和迭代方法。所有子模块都会被正确注册，以便在参数列表中出现。

参数:: modules (list of riemann.Module, 可选) – 用于初始化的模块列表

class riemann.nn.ModuleDict(modules=None)

用于存储模块字典的容器类。

该容器允许使用字符串键存储模块，并提供类似字典的访问方法。所有子模块都会被正确注册。

参数:: modules (dict of {str: riemann.Module}, 可选) – 用于初始化的模块字典

class riemann.nn.ParameterList(parameters=None)

用于存储参数列表的容器类。

该容器允许以列表形式存储多个参数。所有参数都会被正确注册。

参数:: parameters (list of riemann.Parameter, 可选) – 用于初始化的参数列表

class riemann.nn.ParameterDict(parameters=None)

用于存储参数字典的容器类。

该容器允许使用字符串键存储参数，并提供类似字典的访问方法。所有参数都会被正确注册。

参数:: parameters (dict of {str: riemann.Parameter}, 可选) – 用于初始化的参数字典

线性层

class riemann.nn.Linear(in_features, out_features, bias=True)

全连接线性层。

参数:

in_features (int) – 输入特征数量
out_features (int) – 输出特征数量
bias (bool, optional) – 是否包含偏置项

卷积层

class riemann.nn.Conv1d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')

一维卷积层。

对一维输入应用卷积运算，提取特征并生成新的特征图。

参数:

in_channels (int) – 输入的通道数
out_channels (int) – 输出的通道数
kernel_size (int 或 tuple) – 卷积核大小
stride (int 或 tuple, 可选) – 步长
padding (int 或 tuple, 可选) – 输入四周的零填充
dilation (int 或 tuple, 可选) – 核元素之间的间隔
groups (int, 可选) – 输入通道与输出通道之间的分组
bias (bool, 可选) – 是否使用偏置项
padding_mode (str, 可选) – 填充模式

class riemann.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')

二维卷积层。

对二维输入应用卷积运算，提取图像特征并生成新的特征图。

参数:

in_channels (int) – 输入的通道数
out_channels (int) – 输出的通道数
kernel_size (int 或 tuple) – 卷积核大小
stride (int 或 tuple, 可选) – 步长
padding (int 或 tuple, 可选) – 输入四周的零填充
dilation (int 或 tuple, 可选) – 核元素之间的间隔
groups (int, 可选) – 输入通道与输出通道之间的分组
bias (bool, 可选) – 是否使用偏置项
padding_mode (str, 可选) – 填充模式

class riemann.nn.Conv3d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')

三维卷积层。

对三维输入应用卷积运算，常用于视频、体积数据等的特征提取。

参数:

in_channels (int) – 输入的通道数
out_channels (int) – 输出的通道数
kernel_size (int 或 tuple) – 卷积核大小
stride (int 或 tuple, 可选) – 步长
padding (int 或 tuple, 可选) – 输入四周的零填充
dilation (int 或 tuple, 可选) – 核元素之间的间隔
padding_mode (str, 可选) – 填充模式

池化层

class riemann.nn.MaxPool1d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)

一维最大池化层。

对一维输入应用最大池化，用于提取序列数据的关键特征并减少数据维度。

参数:

kernel_size (int 或 tuple) – 池化窗口大小
stride (int 或 tuple, 可选) – 池化窗口移动步长
padding (int 或 tuple, 可选) – 输入四周的零填充
dilation (int 或 tuple, 可选) – 池化窗口元素之间的间隔
return_indices (bool, 可选) – 是否返回最大值的索引
ceil_mode (bool, 可选) – 是否使用向上取整来计算输出形状

class riemann.nn.MaxPool2d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)

二维最大池化层。

对二维输入应用最大池化，常用于图像数据的特征提取和维度减少。

参数:

kernel_size (int 或 tuple) – 池化窗口大小
stride (int 或 tuple, 可选) – 池化窗口移动步长
padding (int 或 tuple, 可选) – 输入四周的零填充
dilation (int 或 tuple, 可选) – 池化窗口元素之间的间隔
return_indices (bool, 可选) – 是否返回最大值的索引
ceil_mode (bool, 可选) – 是否使用向上取整来计算输出形状

class riemann.nn.MaxPool3d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)

三维最大池化层。

对三维输入应用最大池化，常用于视频、体积数据等的特征提取。

参数:

kernel_size (int 或 tuple) – 池化窗口大小
stride (int 或 tuple, 可选) – 池化窗口移动步长
padding (int 或 tuple, 可选) – 输入四周的零填充
dilation (int 或 tuple, 可选) – 池化窗口元素之间的间隔
return_indices (bool, 可选) – 是否返回最大值的索引
ceil_mode (bool, 可选) – 是否使用向上取整来计算输出形状

class riemann.nn.AvgPool1d(kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None)

一维平均池化层。

对一维输入应用平均池化，用于平滑序列数据特征并减少数据维度。

参数:

kernel_size (int 或 tuple) – 池化窗口大小
stride (int 或 tuple, 可选) – 池化窗口移动步长
padding (int 或 tuple, 可选) – 输入四周的零填充
ceil_mode (bool, 可选) – 是否使用向上取整来计算输出形状
count_include_pad (bool, 可选) – 计算平均值时是否包含零填充
divisor_override (int, 可选) – 如果指定，将使用该值作为分母

class riemann.nn.AvgPool2d(kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None)

二维平均池化层。

对二维输入应用平均池化，常用于图像数据的特征平滑和维度减少。

参数:

kernel_size (int 或 tuple) – 池化窗口大小
stride (int 或 tuple, 可选) – 池化窗口移动步长
padding (int 或 tuple, 可选) – 输入四周的零填充
ceil_mode (bool, 可选) – 是否使用向上取整来计算输出形状
count_include_pad (bool, 可选) – 计算平均值时是否包含零填充
divisor_override (int, 可选) – 如果指定，将使用该值作为分母

class riemann.nn.AvgPool3d(kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None)

三维平均池化层。

对三维输入应用平均池化，常用于视频、体积数据等的特征平滑和维度减少。

参数:

kernel_size (int 或 tuple) – 池化窗口大小
stride (int 或 tuple, 可选) – 池化窗口移动步长
padding (int 或 tuple, 可选) – 输入四周的零填充
ceil_mode (bool, 可选) – 是否使用向上取整来计算输出形状
count_include_pad (bool, 可选) – 计算平均值时是否包含零填充
divisor_override (int, 可选) – 如果指定，将使用该值作为分母

自适应池化层

class riemann.nn.AdaptiveAvgPool1d(output_size)

一维自适应平均池化层。

根据指定的输出尺寸自动计算池化核大小和步长，确保输出尺寸始终为指定值。

参数:: output_size (int 或 tuple) – 输出序列长度，可以是整数或 None（表示保持原尺寸）

class riemann.nn.AdaptiveAvgPool2d(output_size)

二维自适应平均池化层。

根据指定的输出尺寸自动计算池化核大小和步长，常用于将任意尺寸的特征图转换为固定尺寸。

参数:: output_size (int 或 tuple) – 输出尺寸，可以是整数或二元组 (H, W)，或 None（表示保持原尺寸）

class riemann.nn.AdaptiveAvgPool3d(output_size)

三维自适应平均池化层。

根据指定的输出尺寸自动计算池化核大小和步长，常用于体积数据的特征提取。

参数:: output_size (int 或 tuple) – 输出尺寸，可以是整数或三元组 (D, H, W)，或 None（表示保持原尺寸）

class riemann.nn.AdaptiveMaxPool1d(output_size, return_indices=False)

一维自适应最大池化层。

根据指定的输出尺寸自动计算池化核大小和步长，确保输出尺寸始终为指定值。

参数:

output_size (int 或 tuple) – 输出序列长度，可以是整数或 None（表示保持原尺寸）
return_indices (bool, 可选) – 是否返回最大值的索引

class riemann.nn.AdaptiveMaxPool2d(output_size, return_indices=False)

二维自适应最大池化层。

根据指定的输出尺寸自动计算池化核大小和步长，常用于将任意尺寸的特征图转换为固定尺寸。

参数:

output_size (int 或 tuple) – 输出尺寸，可以是整数或二元组 (H, W)，或 None（表示保持原尺寸）
return_indices (bool, 可选) – 是否返回最大值的索引

class riemann.nn.AdaptiveMaxPool3d(output_size, return_indices=False)

三维自适应最大池化层。

根据指定的输出尺寸自动计算池化核大小和步长，常用于体积数据的特征提取。

参数:

output_size (int 或 tuple) – 输出尺寸，可以是整数或三元组 (D, H, W)，或 None（表示保持原尺寸）
return_indices (bool, 可选) – 是否返回最大值的索引

归一化层

class riemann.nn.BatchNorm1d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)

一维批量归一化层。

对二维或三维输入张量的通道维度进行归一化，使其具有零均值和单位方差，提升训练收敛性和模型泛化能力。

参数:

num_features (int) – 特征数量（通道维度）
eps (float, 可选) – 避免除零的微小值
momentum (float, 可选) – 运行统计的动量
affine (bool, 可选) – 是否包含可学习的仿射参数
track_running_stats (bool, 可选) – 是否跟踪运行均值和方差

class riemann.nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1)

二维批量归一化层。

参数:

num_features (int) – 特征数量
eps (float, optional) – 避免除零的微小值
momentum (float, optional) – 运行统计的动量

class riemann.nn.BatchNorm3d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)

三维批量归一化层。

对五维输入张量 (N, C, D, H, W) 的通道维度进行归一化，使其具有零均值和单位方差，提升训练收敛性和模型泛化能力。

参数:

num_features (int) – 特征数量（通道维度）
eps (float, 可选) – 避免除零的微小值
momentum (float, 可选) – 运行统计的动量
affine (bool, 可选) – 是否包含可学习的仿射参数
track_running_stats (bool, 可选) – 是否跟踪运行均值和方差

class riemann.nn.LayerNorm(normalized_shape, eps=1e-05, affine=True, device=None, dtype=None)

层归一化层，对指定维度进行归一化处理。

与 torch.nn.LayerNorm 兼容，对输入张量的指定维度进行归一化，使其具有零均值和单位方差。

参数:

normalized_shape (int 或 tuple) – 整数或元组，指定需要归一化的维度
eps (float, optional) – 为避免除零而添加到方差的小值
affine (bool, optional) – 是否包含可学习的仿射参数（gamma 和 beta）
device (optional) – 参数和缓冲区的设备
dtype (optional) – 参数和缓冲区的数据类型

class riemann.nn.Flatten(start_dim=1, end_dim=-1)

展平张量维度的层，移除从 start_dim 到 end_dim 的所有维度。

通常用于卷积层之后、全连接层之前，将多维度的卷积结果展平为一维向量。

参数:

start_dim (int, 可选) – 开始展平的维度
end_dim (int, 可选) – 结束展平的维度

激活函数模块

class riemann.nn.ReLU(inplace=False)

ReLU 激活函数

应用整流线性单元函数逐元素处理输入，ReLU(x) = max(0, x)

参数:: inplace (bool, 可选) – 是否就地操作

class riemann.nn.LeakyReLU(negative_slope=0.01, inplace=False)

Leaky ReLU 激活函数

应用带泄漏的整流线性单元函数逐元素处理输入，LeakyReLU(x) = max(x, negative_slope * x)

参数:

negative_slope (float, 可选) – 负数区域的斜率
inplace (bool, 可选) – 是否就地操作

class riemann.nn.RReLU(lower=0.125, upper=0.3333333333333333, inplace=False)

随机化 Leaky ReLU 激活函数

应用随机化带泄漏的整流线性单元函数逐元素处理输入

参数:

lower (float, 可选) – 均匀分布的下限
upper (float, 可选) – 均匀分布的上限
inplace (bool, 可选) – 是否就地操作

class riemann.nn.PReLU(num_parameters=1, init=0.25)

参数化 ReLU 激活函数

应用参数化的整流线性单元函数逐元素处理输入，a 是可学习的参数

参数:

num_parameters (int, 可选) – 可学习参数的数量
init (float, 可选) – 参数的初始值

class riemann.nn.Sigmoid

Sigmoid 激活函数

应用 Sigmoid 函数逐元素处理输入，将值映射到 [0, 1] 区间

class riemann.nn.Tanh

Tanh 激活函数

应用双曲正切函数逐元素处理输入，将值映射到 [-1, 1] 区间

class riemann.nn.Softmax(dim=None)

Softmax 激活函数

在指定维度上应用 Softmax 函数

参数:: dim (int, 可选) – 应用 softmax 的维度

class riemann.nn.LogSoftmax(dim=None)

Log-Softmax 激活函数

在指定维度上应用 Log-Softmax 函数

参数:: dim (int, 可选) – 应用 log-softmax 的维度

class riemann.nn.GELU

高斯误差线性单元激活函数

应用高斯误差线性单元函数逐元素处理输入，GELU(x) = x * Φ(x)，其中 Φ 是标准正态分布的累积分布函数

class riemann.nn.Softplus(beta=1, threshold=20)

Softplus 激活函数

应用 Softplus 激活函数逐元素处理输入，Softplus(x) = (1 / beta) * log(1 + exp(beta * x))

参数:

beta (float, 可选) – 线性部分的斜率
threshold (float, 可选) – 数值稳定性的阈值

Dropout 层

class riemann.nn.Dropout(p=0.5)

Dropout 层，用于防止过拟合。

参数:: p (float, optional) – 丢弃概率

class riemann.nn.Dropout2d(p=0.5, inplace=False)

2D dropout层，用于防止过拟合。

在训练期间，以概率p随机将输入张量的整个通道置为0，并对剩余通道乘以1/(1-p)进行缩放。在评估期间，不执行任何操作。

参数:

p (float, optional) – 丢弃概率
inplace (bool, optional) – 是否原地操作

class riemann.nn.Dropout3d(p=0.5, inplace=False)

3D dropout层，用于防止过拟合。

在训练期间，以概率p随机将输入张量的整个通道置为0，并对剩余通道乘以1/(1-p)进行缩放。在评估期间，不执行任何操作。

参数:

p (float, optional) – 丢弃概率
inplace (bool, optional) – 是否原地操作

嵌入层

class riemann.nn.Embedding(num_embeddings, embedding_dim, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False, dtype=None, device=None)

嵌入层，将整数索引转换为密集向量。

嵌入层是神经网络中处理分类特征和序列数据的基础组件。

参数:

num_embeddings (int) – 嵌入向量的数量，即词典大小
embedding_dim (int) – 每个嵌入向量的维度
padding_idx (int, 可选) – 如果指定，该索引的嵌入向量不参与梯度计算，且在训练过程中保持不变
max_norm (float, 可选) – 如果指定，所有嵌入向量的范数超过max_norm时，将被重归一化到max_norm
norm_type (float, 可选) – 计算范数时使用的p值，默认为2（L2范数）
scale_grad_by_freq (bool, 可选) – 如果为True，梯度将按mini-batch中每个词的频率进行缩放
sparse (bool, 可选) – 如果为True，权重的梯度将是稀疏张量
dtype (np.dtype, 可选) – 嵌入权重的数据类型
device (str|int|Device, 可选) – 嵌入权重的设备

损失函数模块

class riemann.nn.L1Loss(size_average=None, reduce=None, reduction='mean')

平均绝对误差损失，计算输入与目标值之间的绝对值误差。

参数:

size_average (bool, optional) – 已弃用
reduce (bool, optional) – 已弃用
reduction (str, optional) – 指定输出的聚合方式

class riemann.nn.MSELoss(size_average=None, reduce=None, reduction='mean')

均方误差损失，计算输入与目标值之间的平方误差。

参数:

size_average (bool, optional) – 已弃用
reduce (bool, optional) – 已弃用
reduction (str, optional) – 指定输出的聚合方式

class riemann.nn.NLLLoss(weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean')

负对数似然损失，用于分类任务中的概率预测。

参数:

weight (riemann.TN, optional) – 每个类别的手动缩放权重
size_average (bool, optional) – 已弃用
ignore_index (int, optional) – 指定要忽略的目标值
reduce (bool, optional) – 已弃用
reduction (str, optional) – 指定输出的聚合方式

class riemann.nn.CrossEntropyLoss(weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean')

交叉熵损失，结合了 LogSoftmax 和 NLLLoss 于一个类中，常用于多分类任务。

参数:

weight (riemann.TN, optional) – 每个类别的手动缩放权重
size_average (bool, optional) – 已弃用
ignore_index (int, optional) – 指定要忽略的目标值
reduce (bool, optional) – 已弃用
reduction (str, optional) – 指定输出的聚合方式

class riemann.nn.BCELoss(weight=None, size_average=None, reduce=None, reduction='mean')

二元交叉熵损失，计算目标与输出之间的二元分类误差。

参数:

weight (riemann.TN, optional) – 每个批次元素的手动缩放权重
size_average (bool, optional) – 已弃用
reduce (bool, optional) – 已弃用
reduction (str, optional) – 指定输出的聚合方式

class riemann.nn.BCEWithLogitsLoss(weight=None, size_average=None, reduce=None, reduction='mean')

带logits的二元交叉熵损失，直接对输入logits计算二元交叉熵。

参数:

weight (riemann.TN, optional) – 每个批次元素的手动缩放权重
size_average (bool, optional) – 已弃用
reduce (bool, optional) – 已弃用
reduction (str, optional) – 指定输出的聚合方式

class riemann.nn.HuberLoss(delta=1.0, size_average=None, reduce=None, reduction='mean')

Huber损失函数，当误差小于delta时使用平方误差，否则使用线性误差。

参数:

delta (float, optional) – 损失函数从二次变为线性的阈值
size_average (bool, optional) – 已弃用
reduce (bool, optional) – 已弃用
reduction (str, optional) – 指定输出的聚合方式

class riemann.nn.SmoothL1Loss(beta=1.0, size_average=None, reduce=None, reduction='mean')

平滑L1损失，结合了L1损失和L2损失的优点，在小误差时使用二次损失，大误差时使用线性损失。

参数:

beta (float, optional) – 控制从二次损失到线性损失的转换阈值
size_average (bool, optional) – 已弃用
reduce (bool, optional) – 已弃用
reduction (str, optional) – 指定输出的聚合方式

Transformer模块

class riemann.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None, batch_first=False, device=None, dtype=None)

多头注意力机制，允许模型同时关注来自不同表示子空间的信息。

参数:

embed_dim (int) – 输入和输出向量的维度，必须能被num_heads整除
num_heads (int) – 注意力头的数量
dropout (float, 可选) – 注意力权重的dropout概率
bias (bool, 可选) – 是否在投影层添加偏置
add_bias_kv (bool, 可选) – 是否在key和value序列末尾添加可学习的偏置
add_zero_attn (bool, 可选) – 是否在注意力权重中添加一列零
kdim (int, 可选) – key向量的维度，默认为embed_dim
vdim (int, 可选) – value向量的维度，默认为embed_dim
batch_first (bool, 可选) – 输入输出形状是否为(batch, seq, feature)而非(seq, batch, feature)
device (可选) – 张量设备
dtype (可选) – 张量数据类型

class riemann.nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='relu', layer_norm_eps=1e-05, batch_first=False, norm_first=False, bias=True, device=None, dtype=None)

Transformer编码器的单层，由自注意力机制和前馈网络组成。

参数:

d_model (int) – 输入和输出特征的维度
nhead (int) – 注意力头的数量
dim_feedforward (int, 可选) – 前馈网络隐藏层的维度
dropout (float, 可选) – dropout概率
activation (str, 可选) – 激活函数类型，’relu’或’gelu’
layer_norm_eps (float, 可选) – 层归一化的epsilon值
batch_first (bool, 可选) – 输入输出形状是否为(batch, seq, feature)
norm_first (bool, 可选) – 是否使用Pre-LN模式
bias (bool, 可选) – 是否在线性层添加偏置
device (可选) – 张量设备
dtype (可选) – 张量数据类型

class riemann.nn.TransformerDecoderLayer(d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='relu', layer_norm_eps=1e-05, batch_first=False, norm_first=False, bias=True, device=None, dtype=None)

Transformer解码器的单层，由自注意力、交叉注意力和前馈网络组成。

参数:

d_model (int) – 输入和输出特征的维度
nhead (int) – 注意力头的数量
dim_feedforward (int, 可选) – 前馈网络隐藏层的维度
dropout (float, 可选) – dropout概率
activation (str, 可选) – 激活函数类型，’relu’或’gelu’
layer_norm_eps (float, 可选) – 层归一化的epsilon值
batch_first (bool, 可选) – 输入输出形状是否为(batch, seq, feature)
norm_first (bool, 可选) – 是否使用Pre-LN模式
bias (bool, 可选) – 是否在线性层添加偏置
device (可选) – 张量设备
dtype (可选) – 张量数据类型

class riemann.nn.TransformerEncoder(encoder_layer, num_layers, norm=None, enable_nested_tensor=True, mask_check=True)

由N个堆叠的TransformerEncoderLayer组成的Transformer编码器。

参数:

encoder_layer (TransformerEncoderLayer) – 单个编码器层实例，将被克隆
num_layers (int) – 编码器层的数量
norm (Module, 可选) – 最后的层归一化，可选
enable_nested_tensor (bool, 可选) – 是否启用嵌套张量优化（仅接口兼容）
mask_check (bool, 可选) – 是否进行掩码检查（仅接口兼容）

class riemann.nn.TransformerDecoder(decoder_layer, num_layers, norm=None)

由N个堆叠的TransformerDecoderLayer组成的Transformer解码器。

参数:

decoder_layer (TransformerDecoderLayer) – 单个解码器层实例，将被克隆
num_layers (int) – 解码器层的数量
norm (Module, 可选) – 最后的层归一化，可选

class riemann.nn.Transformer(d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6, dim_feedforward=2048, dropout=0.1, activation='relu', custom_encoder=None, custom_decoder=None, layer_norm_eps=1e-05, batch_first=False, norm_first=False, bias=True, device=None, dtype=None)

包含编码器和解码器的完整Transformer架构。

参数:

d_model (int, 可选) – 编码器/解码器输入的维度
nhead (int, 可选) – 注意力头的数量
num_encoder_layers (int, 可选) – 编码器层的数量
num_decoder_layers (int, 可选) – 解码器层的数量
dim_feedforward (int, 可选) – 前馈网络的维度
dropout (float, 可选) – dropout值
activation (str, 可选) – 激活函数，’relu’或’gelu’
custom_encoder (Module, 可选) – 自定义编码器模块
custom_decoder (Module, 可选) – 自定义解码器模块
layer_norm_eps (float, 可选) – 层归一化的epsilon值
batch_first (bool, 可选) – 输入输出形状是否为(batch, seq, feature)
norm_first (bool, 可选) – 是否在注意力和前馈操作之前执行LayerNorm
bias (bool, 可选) – 线性层和LayerNorm层是否学习加性偏置
device (可选) – 张量设备
dtype (可选) – 张量数据类型

函数式接口

riemann.nn.functional 模块提供了各种神经网络操作的函数式实现。

线性函数

riemann.nn.functional.linear(input, weight, bias=None)

应用线性变换：y = xA^T + b

参数:

input (riemann.TN) – 输入张量，形状为 (*, in_features)
weight (riemann.TN) – 权重张量，形状为 (out_features, in_features)
bias (riemann.TN, optional) – 偏置张量，形状为 (out_features,). 默认值: None

返回:

输出张量，形状为 (*, out_features)

返回类型:

riemann.TN

激活函数

riemann.nn.functional.sigmoid(input)

应用逐元素的 sigmoid 函数：sigmoid(x) = 1 / (1 + exp(-x))

参数:: input (riemann.TN) – 输入张量
返回:: 输出张量
返回类型:: riemann.TN

riemann.nn.functional.silu(input)

应用 Sigmoid Linear Unit (SiLU) 激活函数：silu(x) = x * sigmoid(x)

参数:: input (riemann.TN) – 输入张量
返回:: 输出张量
返回类型:: riemann.TN

riemann.nn.functional.tanh(input)

应用双曲正切激活函数：tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))

参数:: input (riemann.TN) – 输入张量

Dropout函数

riemann.nn.functional.dropout(input, p=0.5, training=True, inplace=False)

在训练期间，以概率p随机将输入张量的元素置为0，并对剩余元素乘以1/(1-p)进行缩放。在评估期间，不执行任何操作。

参数:

input (riemann.TN) – 输入张量
p (float, optional) – 丢弃概率
training (bool, optional) – 是否在训练模式
inplace (bool, optional) – 是否原地操作

返回:

应用dropout后的张量

返回类型:

riemann.TN

riemann.nn.functional.dropout2d(input, p=0.5, training=True, inplace=False)

在训练期间，以概率p随机将输入张量的整个通道置为0，并对剩余通道乘以1/(1-p)进行缩放。在评估期间，不执行任何操作。

参数:

input (riemann.TN) – 输入张量，形状为 (N, C, H, W)
p (float, optional) – 丢弃概率
training (bool, optional) – 是否在训练模式
inplace (bool, optional) – 是否原地操作

返回:

应用dropout后的张量

返回类型:

riemann.TN

riemann.nn.functional.dropout3d(input, p=0.5, training=True, inplace=False)

在训练期间，以概率p随机将输入张量的整个通道置为0，并对剩余通道乘以1/(1-p)进行缩放。在评估期间，不执行任何操作。

参数:

input (riemann.TN) – 输入张量，形状为 (N, C, D, H, W)
p (float, optional) – 丢弃概率
training (bool, optional) – 是否在训练模式
inplace (bool, optional) – 是否原地操作

返回:

应用dropout后的张量

返回类型:

riemann.TN

归一化函数

riemann.nn.functional.batch_norm(input, running_mean=None, running_var=None, weight=None, bias=None, training=False, momentum=0.1, eps=1e-5)

对输入张量应用批量归一化。

参数:

input (riemann.TN) – 输入张量，形状为 (N, C), (N, C, L), (N, C, H, W) 或 (N, C, D, H, W)
running_mean (riemann.TN, optional) – 运行时均值，形状为 (C,)
running_var (riemann.TN, optional) – 运行时方差，形状为 (C,)
weight (riemann.TN, optional) – 可学习的缩放参数γ，形状为 (C,)
bias (riemann.TN, optional) – 可学习的偏移参数β，形状为 (C,)
training (bool, optional) – 是否为训练模式
momentum (float, optional) – 运行时统计量的动量
eps (float, optional) – 数值稳定性的小常数

返回:

归一化后的张量，形状与输入相同

返回类型:

riemann.TN

riemann.nn.functional.layer_norm(input, normalized_shape, weight=None, bias=None, eps=1e-05)

对输入张量的指定维度应用层归一化。

参数:

input (riemann.TN) – 输入张量
normalized_shape (int 或 tuple) – 整数或元组，指定要归一化的维度
weight (riemann.TN, optional) – 可选的权重张量（γ），用于仿射变换
bias (riemann.TN, optional) – 可选的偏置张量（β），用于仿射变换
eps (float, optional) – 为避免除零而添加到方差的小值

返回:

归一化后的张量，形状与输入相同

返回类型:

riemann.TN

嵌入函数

riemann.nn.functional.embedding(input, weight, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False)

从嵌入矩阵中查找输入索引的嵌入向量。

参数:

input (riemann.TN) – 包含索引的张量，形状为任意维度
weight (riemann.TN) – 嵌入矩阵，形状为 (num_embeddings, embedding_dim)
padding_idx (int, optional) – 如果指定，该索引的嵌入向量不参与梯度计算，且在训练过程中保持不变
max_norm (float, optional) – 如果指定，所有嵌入向量的范数超过max_norm时，将被重归一化到max_norm
norm_type (float, optional) – 计算范数时使用的p值，默认为2（L2范数）
scale_grad_by_freq (bool, optional) – 如果为True，梯度将按mini-batch中每个词的频率进行缩放
sparse (bool, optional) – 如果为True，权重的梯度将是稀疏张量

返回:

输出张量，形状为 (*, embedding_dim)，其中``*``是输入的形状

返回类型:

riemann.TN

riemann.nn.functional.softmax(input, dim)

沿指定维度应用 softmax 函数

参数:

input (riemann.TN) – 输入张量
dim (int) – 计算 softmax 的维度

返回:

输出张量

返回类型:

riemann.TN

riemann.nn.functional.log_softmax(input, dim=-1)

应用对数 softmax 函数以获得数值稳定性

参数:

input (riemann.TN) – 输入张量
dim (int, optional) – 计算 log_softmax 的维度

返回:

输出张量

返回类型:

riemann.TN

riemann.nn.functional.relu(input)

应用整流线性单元激活函数：relu(x) = max(0, x)

参数:: input (riemann.TN) – 输入张量
返回:: 输出张量
返回类型:: riemann.TN

riemann.nn.functional.leaky_relu(input, alpha=0.01)

应用带泄漏的整流线性单元激活函数

参数:

input (riemann.TN) – 输入张量
alpha (float, optional) – 负区域的斜率

返回:

输出张量

返回类型:

riemann.TN

riemann.nn.functional.prelu(input, alpha)

应用参数化整流线性单元激活函数

参数:

input (riemann.TN) – 输入张量
alpha (riemann.TN) – 可学习的参数张量

返回:

输出张量

返回类型:

riemann.TN

riemann.nn.functional.rrelu(input, lower=1.0 / 8.0, upper=1.0 / 3.0, training=True)

应用随机化整流线性单元激活函数

参数:

input (riemann.TN) – 输入张量
lower (float, optional) – 均匀分布的下限
upper (float, optional) – 均匀分布的上限
training (bool, optional) – 是否使用随机化的 alpha（训练）或固定的 alpha（评估）

返回:

输出张量

返回类型:

riemann.TN

riemann.nn.functional.gelu(input)

应用高斯误差线性单元激活函数

参数:: input (riemann.TN) – 输入张量
返回:: 输出张量
返回类型:: riemann.TN

riemann.nn.functional.softplus(input, beta=1.0, threshold=20.0)

应用 Softplus 激活函数：softplus(x) = (1 / beta) * log(1 + exp(beta * x))

参数:

input (riemann.TN) – 输入张量
beta (float, optional) – 线性部分的斜率
threshold (float, optional) – 数值稳定性的阈值

返回:

输出张量

返回类型:

riemann.TN

损失函数

riemann.nn.functional.mse_loss(input, target, size_average=None, reduce=None, reduction='mean')

计算均方误差损失

参数:

input (riemann.TN) – 输入张量
target (riemann.TN) – 目标张量
size_average (bool, optional) – 已弃用
reduce (bool, optional) – 已弃用
reduction (str, optional) – 指定输出的聚合方式

返回:

损失值

返回类型:

riemann.TN

riemann.nn.functional.l1_loss(input, target, size_average=None, reduce=None, reduction='mean')

计算 L1（绝对误差）损失

参数:

input (riemann.TN) – 输入张量
target (riemann.TN) – 目标张量
size_average (bool, optional) – 已弃用
reduce (bool, optional) – 已弃用
reduction (str, optional) – 指定输出的聚合方式

返回:

损失值

返回类型:

riemann.TN

riemann.nn.functional.smooth_l1_loss(input, target, size_average=None, reduce=None, reduction='mean', beta=1.0)

计算平滑 L1 损失

参数:

input (riemann.TN) – 输入张量
target (riemann.TN) – 目标张量
size_average (bool, optional) – 已弃用
reduce (bool, optional) – 已弃用
reduction (str, optional) – 指定输出的聚合方式
beta (float, optional) – 损失函数从二次变为线性的阈值

返回:

损失值

返回类型:

riemann.TN

riemann.nn.functional.cross_entropy(input, target, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean', label_smoothing=0.0)

计算交叉熵损失

参数:

input (riemann.TN) – 输入张量
target (riemann.TN) – 目标张量
weight (riemann.TN, optional) – 每个类别的手动缩放权重
size_average (bool, optional) – 已弃用
ignore_index (int, optional) – 指定要忽略的目标值
reduce (bool, optional) – 已弃用
reduction (str, optional) – 指定输出的聚合方式
label_smoothing (float, optional) – 标签平滑的量

返回:

损失值

返回类型:

riemann.TN

riemann.nn.functional.binary_cross_entropy_with_logits(input, target, weight=None, size_average=None, reduce=None, reduction='mean', pos_weight=None)

计算带 logits 的二元交叉熵损失

参数:

input (riemann.TN) – 输入张量
target (riemann.TN) – 目标张量
weight (riemann.TN, optional) – 每个批次元素的手动缩放权重
size_average (bool, optional) – 已弃用
reduce (bool, optional) – 已弃用
reduction (str, optional) – 指定输出的聚合方式
pos_weight (riemann.TN, optional) – 正类的权重

返回:

损失值

返回类型:

riemann.TN

riemann.nn.functional.huber_loss(input, target, delta=1.0, size_average=None, reduce=None, reduction='mean')

计算 Huber 损失

参数:

input (riemann.TN) – 输入张量
target (riemann.TN) – 目标张量
delta (float, optional) – 损失函数从二次变为线性的阈值
size_average (bool, optional) – 已弃用
reduce (bool, optional) – 已弃用
reduction (str, optional) – 指定输出的聚合方式

返回:

损失值

返回类型:

riemann.TN

riemann.nn.functional.nll_loss(input, target, weight=None, ignore_index=-100, reduction='mean')

计算负对数似然损失

参数:

input (riemann.TN) – 输入张量
target (riemann.TN) – 目标张量
weight (riemann.TN, optional) – 每个类别的手动缩放权重
ignore_index (int, optional) – 指定要忽略的目标值
reduction (str, optional) – 指定输出的聚合方式

返回:

损失值

返回类型:

riemann.TN

卷积函数

riemann.nn.functional.conv1d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1)

对输入信号应用 1D 卷积

参数:

input (riemann.TN) – 输入张量，形状为 (N, C_in, L_in)
weight (riemann.TN) – 权重张量，形状为 (C_out, C_in/groups, K)
bias (riemann.TN, optional) – 偏置张量，形状为 (C_out). 默认值: None
stride (int or tuple, optional) – 卷积步长
padding (int or tuple, optional) – 添加到输入两侧的零填充
dilation (int or tuple, optional) – 内核元素之间的间距
groups (int, optional) – 从输入通道到输出通道的阻塞连接数

返回:

输出张量，形状为 (N, C_out, L_out)

返回类型:

riemann.TN

riemann.nn.functional.conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1)

对输入图像应用 2D 卷积

参数:

input (riemann.TN) – 输入张量，形状为 (N, C_in, H_in, W_in)
weight (riemann.TN) – 权重张量，形状为 (C_out, C_in/groups, K_h, K_w)
bias (riemann.TN, optional) – 偏置张量，形状为 (C_out). 默认值: None
stride (int or tuple, optional) – 卷积步长
padding (int or tuple, optional) – 添加到输入两侧的零填充
dilation (int or tuple, optional) – 内核元素之间的间距
groups (int, optional) – 从输入通道到输出通道的阻塞连接数

返回:

输出张量，形状为 (N, C_out, H_out, W_out)

返回类型:

riemann.TN

riemann.nn.functional.conv3d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1)

对输入体积应用 3D 卷积

参数:

input (riemann.TN) – 输入张量，形状为 (N, C_in, D_in, H_in, W_in)
weight (riemann.TN) – 权重张量，形状为 (C_out, C_in/groups, K_d, K_h, K_w)
bias (riemann.TN, optional) – 偏置张量，形状为 (C_out). 默认值: None
stride (int or tuple, optional) – 卷积步长
padding (int or tuple, optional) – 添加到输入所有侧面的零填充
dilation (int or tuple, optional) – 内核元素之间的间距
groups (int, optional) – 从输入通道到输出通道的阻塞连接数

返回:

输出张量，形状为 (N, C_out, D_out, H_out, W_out)

返回类型:

riemann.TN

池化函数

riemann.nn.functional.max_pool1d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False)

对输入信号应用 1D 最大池化

参数:

input (riemann.TN) – 输入张量，形状为 (N, C, L_in)
kernel_size (int or tuple) – 池化窗口大小
stride (int or tuple, optional) – 池化窗口移动步长。默认值: kernel_size
padding (int or tuple, optional) – 添加到输入两侧的零填充
dilation (int or tuple, optional) – 池化窗口元素之间的间距
ceil_mode (bool, optional) – 是否使用向上取整来计算输出形状
return_indices (bool, optional) – 是否返回最大值的索引

返回:

输出张量，形状为 (N, C, L_out)，如果 return_indices 为 True，则返回元组 (TN, TN)

返回类型:

riemann.TN or tuple

riemann.nn.functional.max_pool2d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False)

对输入图像应用 2D 最大池化

参数:

input (riemann.TN) – 输入张量，形状为 (N, C, H_in, W_in)
kernel_size (int or tuple) – 池化窗口大小
stride (int or tuple, optional) – 池化窗口移动步长。默认值: kernel_size
padding (int or tuple, optional) – 添加到输入两侧的零填充
dilation (int or tuple, optional) – 池化窗口元素之间的间距
ceil_mode (bool, optional) – 是否使用向上取整来计算输出形状
return_indices (bool, optional) – 是否返回最大值的索引

返回:

输出张量，形状为 (N, C, H_out, W_out)，如果 return_indices 为 True，则返回元组 (TN, TN)

返回类型:

riemann.TN or tuple

riemann.nn.functional.max_pool3d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False)

对输入体积数据应用 3D 最大池化

参数:

input (riemann.TN) – 输入张量，形状为 (N, C, D_in, H_in, W_in)
kernel_size (int or tuple) – 池化窗口大小
stride (int or tuple, optional) – 池化窗口移动步长。默认值: kernel_size
padding (int or tuple, optional) – 添加到输入所有侧面的零填充
dilation (int or tuple, optional) – 池化窗口元素之间的间距
ceil_mode (bool, optional) – 是否使用向上取整来计算输出形状
return_indices (bool, optional) – 是否返回最大值的索引

返回:

输出张量，形状为 (N, C, D_out, H_out, W_out)，如果 return_indices 为 True，则返回元组 (TN, TN)

返回类型:

riemann.TN or tuple

riemann.nn.functional.avg_pool1d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None)

对输入信号应用 1D 平均池化

参数:

input (riemann.TN) – 输入张量，形状为 (N, C, L_in)
kernel_size (int or tuple) – 池化窗口大小
stride (int or tuple, optional) – 池化窗口移动步长。默认值: kernel_size
padding (int or tuple, optional) – 添加到输入两侧的零填充
ceil_mode (bool, optional) – 是否使用向上取整来计算输出形状
count_include_pad (bool, optional) – 计算平均值时是否包含零填充
divisor_override (int, optional) – 如果指定，将使用该值作为分母

返回:

输出张量，形状为 (N, C, L_out)

返回类型:

riemann.TN

riemann.nn.functional.avg_pool2d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None)

对输入图像应用 2D 平均池化

参数:

input (riemann.TN) – 输入张量，形状为 (N, C, H_in, W_in)
kernel_size (int or tuple) – 池化窗口大小
stride (int or tuple, optional) – 池化窗口移动步长。默认值: kernel_size
padding (int or tuple, optional) – 添加到输入两侧的零填充
ceil_mode (bool, optional) – 是否使用向上取整来计算输出形状
count_include_pad (bool, optional) – 计算平均值时是否包含零填充
divisor_override (int, optional) – 如果指定，将使用该值作为分母

返回:

输出张量，形状为 (N, C, H_out, W_out)

返回类型:

riemann.TN

riemann.nn.functional.avg_pool3d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None)

对输入体积数据应用 3D 平均池化

参数:

input (riemann.TN) – 输入张量，形状为 (N, C, D_in, H_in, W_in)
kernel_size (int or tuple) – 池化窗口大小
stride (int or tuple, optional) – 池化窗口移动步长。默认值: kernel_size
padding (int or tuple, optional) – 添加到输入所有侧面的零填充
ceil_mode (bool, optional) – 是否使用向上取整来计算输出形状
count_include_pad (bool, optional) – 计算平均值时是否包含零填充
divisor_override (int, optional) – 如果指定，将使用该值作为分母

返回:

输出张量，形状为 (N, C, D_out, H_out, W_out)

返回类型:

riemann.TN

工具函数

riemann.nn.functional.one_hot(target, num_classes)

将类别索引转换为 one-hot 编码张量

参数:

target (riemann.TN) – 目标张量，形状为 (N, *)
num_classes (int) – 类别数量

返回:

One-hot 编码张量，形状为 (N, *, num_classes)

返回类型:

riemann.TN

riemann.nn.functional.unfold(input, kernel_size, dilation=1, padding=0, stride=1)

从批处理输入张量中提取滑动局部块

参数:

input (riemann.TN) – 输入张量，形状为 (N, C, H, W)
kernel_size (int or tuple) – 滑动块大小
dilation (int or tuple, optional) – 内核元素之间的间距
padding (int or tuple, optional) – 添加到输入两侧的零填充
stride (int or tuple, optional) – 滑动块的步长

返回:

展开后的张量，形状为 (N, C * kernel_size[0] * kernel_size[1], L)

返回类型:

riemann.TN

riemann.nn.functional.fold(input, output_size, kernel_size, dilation=1, padding=0, stride=1)

将展开的张量折叠回原始形状

参数:

input (riemann.TN) – 输入张量，形状为 (N, C * kernel_size[0] * kernel_size[1], L)
output_size (int or tuple) – 输出张量大小 (H, W)
kernel_size (int or tuple) – 滑动块大小
dilation (int or tuple, optional) – 内核元素之间的间距
padding (int or tuple, optional) – 添加到输入两侧的零填充
stride (int or tuple, optional) – 滑动块的步长

返回:

折叠后的张量，形状为 (N, C, H, W)

返回类型:

riemann.TN

riemann.nn.functional.unfold2d(input, kernel_size, dilation=1, padding=0, stride=1)

从2D输入张量中提取滑动局部块（unfold的2D专用版本）

参数:

input (riemann.TN) – 输入张量，形状为 (N, C, H, W)
kernel_size (int or tuple) – 滑动块大小
dilation (int or tuple, optional) – 内核元素之间的间距
padding (int or tuple, optional) – 添加到输入两侧的零填充
stride (int or tuple, optional) – 滑动块的步长

返回:

展开后的张量，形状为 (N, C * kernel_size[0] * kernel_size[1], L)

返回类型:

riemann.TN

riemann.nn.functional.unfold3d(input, kernel_size, dilation=1, padding=0, stride=1)

从3D输入张量中提取滑动局部块（unfold的3D专用版本）

参数:

input (riemann.TN) – 输入张量，形状为 (N, C, D, H, W)
kernel_size (int or tuple) – 滑动块大小
dilation (int or tuple, optional) – 内核元素之间的间距
padding (int or tuple, optional) – 添加到输入所有侧面的零填充
stride (int or tuple, optional) – 滑动块的步长

返回:

展开后的张量，形状为 (N, C * kernel_size[0] * kernel_size[1] * kernel_size[2], L)

返回类型:

riemann.TN

数据集

数据集类

class riemann.utils.Dataset

抽象数据集基类，定义了所有数据集必须实现的标准接口。

__len__(): 返回数据集中的样本数量。

__getitem__(index): 根据给定索引获取数据集中的单个样本。

class riemann.utils.TensorDataset(*tensors)

简单的张量数据集实现，将多个张量的第一个维度作为数据集维度。

参数:: *tensors (riemann.TN) – 可变数量的张量，所有张量的第一个维度大小必须相同

__len__(): 返回数据集的大小，即张量的第一个维度大小。

__getitem__(index): 获取指定索引处的样本数据。

数据加载器

class riemann.utils.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0, collate_fn=None, drop_last=False)

高效的数据加载器，支持批次处理、数据洗牌和多进程加载。

参数:

dataset (riemann.utils.Dataset) – 要加载数据的数据集
batch_size (int, optional) – 每个批次的大小，默认为1
shuffle (bool, optional) – 是否在每个 epoch 开始时洗牌数据，默认为False
num_workers (int, optional) – 数据加载的工作进程数，0表示主进程加载，默认为0
collate_fn (callable, optional) – 批次处理函数，用于将样本组合成批次，默认使用default_collate
drop_last (bool, optional) – 如果数据集大小不能被批次大小整除，是否丢弃最后一个不完整的批次，默认为False

__len__(): 返回数据加载器的批次数目。

__iter__(): 返回数据加载器的迭代器。

数据集工具函数

riemann.utils.default_collate(batch)

默认的批次处理函数，将一批样本数据转换为适合模型输入的张量格式。

参数:: batch (list) – 一个批次的样本列表，每个样本可以是各种数据类型。
返回:: 根据输入类型组合成的批次数据。

riemann.utils.clip_grad_norm_(parameters, max_norm, norm_type=2.0, error_if_nonfinite=False)

按范数裁剪梯度。

参数:

parameters (Iterable[riemann.TN]) – 需要裁剪梯度的参数集合
max_norm (float or int) – 梯度的最大范数
norm_type (float or int, optional) – 范数类型，默认为2（L2范数）
error_if_nonfinite (bool, optional) – 如果梯度包含非有限值（如NaN或inf），是否抛出错误，默认为False

返回:

裁剪前的梯度范数

返回类型:

float

riemann.utils.clip_grad_value_(parameters, clip_value, error_if_nonfinite=False)

按值裁剪梯度。

参数:

parameters (Iterable[riemann.TN]) – 需要裁剪梯度的参数集合
clip_value (float or int) – 梯度裁剪的阈值
error_if_nonfinite (bool, optional) – 如果梯度包含非有限值（如NaN或inf），是否抛出错误，默认为False

视觉

数据集

class riemann.vision.datasets.MNIST(root, train=True, transform=None, target_transform=None)

MNIST数据集类，用于加载和处理MNIST手写数字数据集。

参数:

root (str) – 数据集的根目录
train (bool) – 是否加载训练集，默认为True
transform (callable, optional) – 应用于图像的变换函数
target_transform (callable, optional) – 应用于目标的变换函数

__len__(): 返回数据集中的样本数量。

__getitem__(index): 根据给定索引获取数据集中的单个样本。

class riemann.vision.datasets.EasyMNIST(root, train=True, onehot_label=True, download=False)

继承自MNIST的子类，在初始化时对图像数据应用归一化、标准化、展开转换，对标签作onehot编码或转换为标量张量。

参数:

root (str) – 数据集的根目录
train (bool) – 是否加载训练集，默认为True
onehot_label (bool) – 是否使用one-hot编码的标签，默认为True
download (bool, optional) – 如果数据集不存在是否下载，默认为False

__len__(): 返回数据集的大小。

__getitem__(index): 获取指定索引处的样本数据。

class riemann.vision.datasets.FashionMNIST(root, train=True, transform=None, target_transform=None, download=False)

Fashion-MNIST数据集类，用于加载和处理Fashion-MNIST时尚产品数据集。

参数:

root (str) – 数据集的根目录
train (bool) – 是否加载训练集，默认为True
transform (callable, optional) – 应用于图像的变换函数
target_transform (callable, optional) – 应用于目标的变换函数
download (bool, optional) – 如果数据集不存在是否下载，默认为False

classes: 类别名称列表：[‘T-shirt/top’, ‘Trouser’, ‘Pullover’, ‘Dress’, ‘Coat’, ‘Sandal’, ‘Shirt’, ‘Sneaker’, ‘Bag’, ‘Ankle boot’]

__len__(): 返回数据集中的样本数量。

__getitem__(index): 根据给定索引获取数据集中的单个样本。

class riemann.vision.datasets.CIFAR10(root, train=True, transform=None, target_transform=None, download=False)

CIFAR-10数据集类，用于加载和处理CIFAR-10图像数据集。

参数:

root (str) – 数据集的根目录
train (bool) – 是否加载训练集，默认为True
transform (callable, optional) – 应用于图像的变换函数
target_transform (callable, optional) – 应用于目标的变换函数
download (bool, optional) – 如果数据集不存在是否下载，默认为False

__len__(): 返回数据集的大小。

__getitem__(index): 获取指定索引处的样本数据。

class riemann.vision.datasets.Flowers102(root, split='train', transform=None, target_transform=None, download=False)

Oxford 102 Flower数据集类，用于加载和处理花卉分类数据集。

参数:

root (str) – 数据集的根目录
split (str, optional) – 数据集划分（’train’、’val’ 或 ‘test’），默认为’train’
transform (callable, optional) – 应用于图像的变换函数
target_transform (callable, optional) – 应用于目标的变换函数
download (bool, optional) – 如果数据集不存在是否下载，默认为False

__len__(): 返回数据集的大小。

__getitem__(index): 获取指定索引处的样本数据。

class riemann.vision.datasets.OxfordIIITPet(root, split='trainval', target_types='category', transform=None, target_transform=None, download=False)

Oxford-IIIT Pet数据集类，用于加载和处理宠物分类数据集。

参数:

root (str) – 数据集的根目录
split (str, optional) – 数据集划分（’trainval’ 或 ‘test’），默认为’trainval’
target_types (str 或 list, optional) – 目标类型（’category’、’binary-category’ 或 ‘segmentation’），默认为’category’
transform (callable, optional) – 应用于图像的变换函数
target_transform (callable, optional) – 应用于目标的变换函数
download (bool, optional) – 如果数据集不存在是否下载，默认为False

__len__(): 返回数据集的大小。

__getitem__(index): 获取指定索引处的样本数据。

class riemann.vision.datasets.LFWPeople(root, split='10fold', image_set='funneled', transform=None, target_transform=None, download=False)

LFW People数据集类，用于加载和处理人脸识别数据集。

参数:

root (str) – 数据集的根目录
split (str, optional) – 数据集划分（’10fold’、’train’ 或 ‘test’），默认为’10fold’
image_set (str, optional) – 图像对齐类型（’original’、’funneled’ 或 ‘deepfunneled’），默认为’funneled’
transform (callable, optional) – 应用于图像的变换函数
target_transform (callable, optional) – 应用于目标的变换函数
download (bool, optional) – 如果数据集不存在是否下载，默认为False

classes: 人物名称列表。

__len__(): 返回数据集的大小。

__getitem__(index): 获取指定索引处的样本数据。

class riemann.vision.datasets.SVHN(root, split='train', transform=None, target_transform=None, download=False)

SVHN（街景门牌号码）数据集类，用于加载和处理数字识别数据集。

参数:

root (str) – 数据集的根目录
split (str, optional) – 数据集划分（’train’、’test’ 或 ‘extra’），默认为’train’
transform (callable, optional) – 应用于图像的变换函数
target_transform (callable, optional) – 应用于目标的变换函数
download (bool, optional) – 如果数据集不存在是否下载，默认为False

__len__(): 返回数据集的大小。

__getitem__(index): 获取指定索引处的样本数据。

class riemann.vision.datasets.DatasetFolder(root, loader, extensions=None, transform=None, target_transform=None, is_valid_file=None, allow_empty=False)

通用文件夹数据集类，用于从文件夹加载自定义数据集。

参数:

root (str) – 数据集的根目录路径
loader (callable) – 图像加载函数
extensions (tuple, optional) – 允许的文件扩展名元组
transform (callable, optional) – 应用于图像的变换函数
target_transform (callable, optional) – 应用于目标的变换函数
is_valid_file (callable, optional) – 验证文件是否有效的函数
allow_empty (bool) – 是否允许空文件夹，默认为False

classes: 类别名称列表。

class_to_idx: 类别名称到索引的映射字典。

__len__(): 返回数据集中的样本数量。

__getitem__(index): 根据给定索引获取数据集中的单个样本。

class riemann.vision.datasets.ImageFolder(root, transform=None, target_transform=None, loader=None, is_valid_file=None)

图像文件夹数据集类，继承自 DatasetFolder，用于从文件夹加载图像数据集。

参数:

root (str) – 数据集的根目录路径
transform (callable, optional) – 应用于图像的变换函数
target_transform (callable, optional) – 应用于目标的变换函数
loader (callable, optional) – 图像加载函数，默认为 PIL Image 加载
is_valid_file (callable, optional) – 验证文件是否有效的函数

classes: 类别名称列表。

class_to_idx: 类别名称到索引的映射字典。

__len__(): 返回数据集中的样本数量。

__getitem__(index): 根据给定索引获取数据集中的单个样本。

图像变换

class riemann.vision.transforms.Transform

所有变换类的基类。

__call__(img): 执行变换。

class riemann.vision.transforms.Compose(transforms)

将多个变换组合成一个变换。

参数:: transforms (list of Transform objects) – 要组合的变换列表

class riemann.vision.transforms.ToTensor: 将PIL图像或NumPy数组转换为TN张量。

class riemann.vision.transforms.ToPILImage: 将TN张量或NumPy数组转换为PIL图像。

class riemann.vision.transforms.Normalize(mean, std, inplace=False)

使用均值和标准差标准化张量。

参数:

mean (sequence) – 每个通道的均值
std (sequence) – 每个通道的标准差
inplace (bool, optional) – 是否原地操作，默认为False

class riemann.vision.transforms.Resize(size, interpolation=BILINEAR)

调整PIL图像大小。

参数:

size (int or tuple) – 目标大小。如果是int，则较小边会被调整为该大小，保持宽高比。如果是(h, w)，则直接调整为该大小。
interpolation (int, optional) – 插值方法，默认为BILINEAR

class riemann.vision.transforms.CenterCrop(size)

中心裁剪。

参数:: size (int or tuple) – 裁剪大小。如果是int，则裁剪为正方形(size, size)。如果是(h, w)，则裁剪为该大小。

class riemann.vision.transforms.RandomHorizontalFlip(p=0.5)

随机水平翻转。

参数:: p (float, optional) – 翻转概率，默认为0.5

class riemann.vision.transforms.RandomVerticalFlip(p=0.5)

随机垂直翻转。

参数:: p (float, optional) – 翻转概率，默认为0.5

class riemann.vision.transforms.RandomRotation(degrees, resample=NEAREST, expand=False, center=None)

随机旋转。

参数:

degrees (int or tuple) – 旋转角度范围。如果是int，则在(-degrees, degrees)范围内选择。如果是(min, max)，则在(min, max)范围内选择。
resample (int, optional) – 重采样方法，默认为NEAREST
expand (bool, optional) – 是否扩展图像以适应旋转，默认为False
center (tuple, optional) – 旋转中心，默认为图像中心

class riemann.vision.transforms.ColorJitter(brightness=0, contrast=0, saturation=0, hue=0)

随机颜色变换。

参数:

brightness (float or tuple) – 亮度调整因子
contrast (float or tuple) – 对比度调整因子
saturation (float or tuple) – 饱和度调整因子
hue (float or tuple) – 色调调整因子

class riemann.vision.transforms.Grayscale(num_output_channels=1)

将图像转换为灰度。

参数:: num_output_channels (int) – 输出通道数，1或3，默认为1

class riemann.vision.transforms.RandomGrayscale(p=0.1)

随机转换为灰度。

参数:: p (float, optional) – 转换为灰度的概率，默认为0.1

class riemann.vision.transforms.RandomCrop(size, padding=None)

在随机位置裁剪图像。

参数:

size (int or tuple) – 裁剪大小。如果是int，则裁剪为正方形(size, size)。如果是(h, w)，则裁剪为该大小。
padding (int or tuple, optional) – 填充大小，默认为None

class riemann.vision.transforms.RandomResizedCrop(size, scale=(0.08, 1.0), ratio=(3 / 4, 4 / 3), interpolation=BILINEAR)

随机裁剪并调整大小。

参数:

size (int or tuple) – 目标大小。如果是int，则调整为正方形(size, size)。如果是(h, w)，则调整为该大小。
scale (tuple, optional) – 裁剪面积相对于原图的比例范围，默认为(0.08, 1.0)
ratio (tuple, optional) – 裁剪的宽高比范围，默认为(3/4, 4/3)
interpolation (int, optional) – 插值方法，默认为BILINEAR

class riemann.vision.transforms.FiveCrop(size)

五裁剪。

参数:: size (int or tuple) – 裁剪大小。如果是int，则裁剪为正方形(size, size)。如果是(h, w)，则裁剪为该大小。

class riemann.vision.transforms.TenCrop(size, vertical_flip=False)

十裁剪。

参数:

size (int or tuple) – 裁剪大小。如果是int，则裁剪为正方形(size, size)。如果是(h, w)，则裁剪为该大小。
vertical_flip (bool, optional) – 是否包括垂直翻转版本，默认为False

class riemann.vision.transforms.Pad(padding, fill=0, padding_mode='constant')

填充。

参数:

padding (int or tuple) – 填充大小。如果是int，则在所有方向填充相同大小。如果是(pad_l, pad_r, pad_t, pad_b)，则分别指定左右上下的填充大小。如果是(pad_h, pad_w)，则分别指定高度和宽度方向的填充大小。
fill (int or tuple) – 填充值，默认为0
padding_mode (str, optional) – 填充模式，默认为’constant’

class riemann.vision.transforms.Lambda(lambd)

使用用户定义的lambda函数作为变换。

参数:: lambd (function) – Lambda函数

class riemann.vision.transforms.PILToTensor: 将PIL Image转换为张量（不缩放）。

class riemann.vision.transforms.ConvertImageDtype(dtype)

转换图像数据类型。

参数:: dtype (torch.dtype) – 目标数据类型

class riemann.vision.transforms.GaussianBlur(kernel_size, sigma=(0.1, 2.0))

对图像应用高斯模糊。

参数:

kernel_size (int or tuple) – 高斯核的大小
sigma (tuple, optional) – 高斯核的标准差范围

class riemann.vision.transforms.RandomAffine(degrees, translate=None, scale=None, shear=None, resample=NEAREST, fillcolor=0)

随机仿射变换。

参数:

degrees (float or tuple) – 旋转角度范围
translate (tuple, optional) – 平移范围
scale (tuple, optional) – 缩放范围
shear (float or tuple, optional) – 剪切角度范围
resample (int, optional) – 重采样模式
fillcolor (int, optional) – 填充颜色

class riemann.vision.transforms.RandomPerspective(distortion_scale=0.5, p=0.5, interpolation=BILINEAR, fill=0)

随机透视变换。

参数:

distortion_scale (float) – 失真程度
p (float) – 应用变换的概率
interpolation (int) – 插值模式
fill (int or tuple) – 填充值

class riemann.vision.transforms.RandomErasing(p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0, inplace=False)

随机擦除，用于数据增强。

参数:

p (float) – 应用擦除的概率
scale (tuple) – 擦除区域面积范围
ratio (tuple) – 擦除区域长宽比范围
value (int or float or tuple) – 擦除填充值
inplace (bool) – 是否原地操作

class riemann.vision.transforms.AutoAugment(policy=AutoAugmentPolicy.IMAGENET)

自动数据增强，基于学习策略。

参数:: policy (AutoAugmentPolicy) – 增强策略

class riemann.vision.transforms.RandAugment(num_ops=2, magnitude=9, num_magnitude_bins=31, interpolation=BILINEAR, fill=None)

随机数据增强。

参数:

num_ops (int) – 操作数量
magnitude (int) – 增强幅度
num_magnitude_bins (int) – 幅度分箱数
interpolation (int) – 插值模式
fill (int or tuple or None) – 填充值

class riemann.vision.transforms.TrivialAugmentWide(num_magnitude_bins=31, interpolation=BILINEAR, fill=None)

宽范围简单增强。

参数:

num_magnitude_bins (int) – 幅度分箱数
interpolation (int) – 插值模式
fill (int or tuple or None) – 填充值

class riemann.vision.transforms.SanitizeBoundingBox(labels_format='xyxy', min_size=1)

边界框清理。

参数:

labels_format (str) – 边界框格式
min_size (int) – 最小尺寸

class riemann.vision.transforms.Invert: 颜色反转。

class riemann.vision.transforms.Posterize(bits)

减少颜色位数。

参数:: bits (int) – 保留的位数

class riemann.vision.transforms.Solarize(threshold)

反转高于阈值的像素。

参数:: threshold (int) – 阈值

class riemann.vision.transforms.Equalize: 直方图均衡化。

class riemann.vision.transforms.AutoContrast: 自动对比度调整。

class riemann.vision.transforms.Sharpness(sharpness_factor)

锐度调整。

参数:: sharpness_factor (float) – 锐度因子

class riemann.vision.transforms.Brightness(brightness_factor)

亮度调整。

参数:: brightness_factor (float) – 亮度因子

class riemann.vision.transforms.Contrast(contrast_factor)

对比度调整。

参数:: contrast_factor (float) – 对比度因子

class riemann.vision.transforms.Saturation(saturation_factor)

饱和度调整。

参数:: saturation_factor (float) – 饱和度因子

class riemann.vision.transforms.Hue(hue_factor)

色调调整。

参数:: hue_factor (float) – 色调因子

优化

优化器

class riemann.optim.Optimizer(params, defaults)

所有优化器的基类。

参数:

params (Iterable[riemann.TN or riemann.nn.Parameter] or List[Dict[str, Any]]) – 待优化参数的迭代器或定义参数组的字典列表
defaults (Dict[str, Any]) – 优化器的默认超参数

step(closure=None)

执行单个优化步骤

参数:: closure (callable, optional) – 重新评估模型并返回损失的闭包
返回:: 如果提供了closure，则返回损失值，否则返回None
返回类型:: float or None

zero_grad(set_to_none=False)

将所有参数的梯度设置为零

参数:: set_to_none (bool, optional) – 是否将梯度设置为None而不是零

add_param_group(param_group)

向优化器添加参数组

参数:: param_group (Dict[str, Any]) – 要添加的参数组

state_dict()

返回优化器的状态字典

返回:: 优化器状态
返回类型:: Dict[str, Any]

load_state_dict(state_dict)

加载优化器状态

参数:: state_dict (Dict[str, Any]) – 要加载的状态字典

class riemann.optim.GD(params, lr=0.01, weight_decay=0.0)

梯度下降（Gradient Descent）优化器

参数:

params (Iterable[riemann.TN or riemann.nn.Parameter] or List[Dict[str, Any]]) – 待优化参数的迭代器或定义参数组的字典列表
lr (float, optional) – 学习率
weight_decay (float, optional) – 权重衰减（L2正则化）系数

step(): 执行单个优化步骤

class riemann.optim.SGD(params, lr=0.01, momentum=0.0, weight_decay=0.0, dampening=0.0, nesterov=False)

随机梯度下降（Stochastic Gradient Descent）优化器

参数:

params (Iterable[riemann.TN or riemann.nn.Parameter] or List[Dict[str, Any]]) – 待优化参数的迭代器或定义参数组的字典列表
lr (float, optional) – 学习率
momentum (float, optional) – 动量因子
weight_decay (float, optional) – 权重衰减（L2正则化）系数
dampening (float, optional) – 动量抑制系数
nesterov (bool, optional) – 是否启用Nesterov动量

step(closure=None)

执行单个优化步骤

参数:: closure (callable, optional) – 重新评估模型并返回损失的闭包
返回:: 如果提供了closure，则返回损失值，否则返回None
返回类型:: float or None

class riemann.optim.Adam(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False)

Adam（Adaptive Moment Estimation）优化器

参数:

params (Iterable[riemann.TN or riemann.nn.Parameter] or List[Dict[str, Any]]) – 待优化参数的迭代器或定义参数组的字典列表
lr (float, optional) – 学习率
betas (Tuple[float, float], optional) – 用于计算梯度及其平方的运行平均值的系数
eps (float, optional) – 添加到分母以提高数值稳定性的项
weight_decay (float, optional) – 权重衰减（L2正则化）系数
amsgrad (bool, optional) – 是否使用AMSGrad变体

step(): 执行单个优化步骤

class riemann.optim.Adagrad(params, lr=0.01, lr_decay=0.0, weight_decay=0.0, initial_accumulator_value=0.0, eps=1e-10)

Adagrad（Adaptive Gradient Algorithm）优化器

参数:

params (Iterable[riemann.TN or riemann.nn.Parameter] or List[Dict[str, Any]]) – 待优化参数的迭代器或定义参数组的字典列表
lr (float, optional) – 学习率
lr_decay (float, optional) – 学习率衰减
weight_decay (float, optional) – 权重衰减（L2正则化）系数
initial_accumulator_value (float, optional) – 累加器的初始值
eps (float, optional) – 添加到分母以提高数值稳定性的项

step(): 执行单个优化步骤

class riemann.optim.LBFGS(params, lr=1.0, max_iter=20, max_eval=None, tolerance_grad=1e-05, tolerance_change=1e-09, history_size=100, line_search_fn=None)

L-BFGS（Limited-memory Broyden-Fletcher-Goldfarb-Shanno）优化器

参数:

params (Iterable[riemann.TN or riemann.nn.Parameter] or List[Dict[str, Any]]) – 待优化参数的迭代器或定义参数组的字典列表
lr (float, optional) – 学习率
max_iter (int, optional) – 每个优化步骤的最大迭代次数
max_eval (int, optional) – 每个优化步骤的最大函数评估次数
tolerance_grad (float, optional) – 梯度收敛阈值
tolerance_change (float, optional) – 参数变化收敛阈值
history_size (int, optional) – 历史缓冲区大小
line_search_fn (callable, optional) – 线搜索函数

step(closure)

执行单个优化步骤

参数:: closure (callable) – 重新评估模型并返回损失的闭包
返回:: 损失值
返回类型:: float

class riemann.optim.AdamW(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01, amsgrad=False)

AdamW（Adam with Weight Decay）优化器

Adam的改进版本，将权重衰减作为独立的正则化项处理，而非Adam中的梯度修改。这使得权重衰减能够更有效地作为L2正则化，避免了Adam中原有的权重衰减副作用。

参数:

params (Iterable[riemann.TN or riemann.nn.Parameter] or List[Dict[str, Any]]) – 待优化参数的迭代器或定义参数组的字典列表
lr (float, optional) – 学习率
betas (Tuple[float, float], optional) – 用于计算梯度及其平方的运行平均值的系数
eps (float, optional) – 添加到分母以提高数值稳定性的项
weight_decay (float, optional) – 权重衰减（L2正则化）系数
amsgrad (bool, optional) – 是否使用AMSGrad变体

step(): 执行单个优化步骤

class riemann.optim.RMSprop(params, lr=1e-2, alpha=0.99, eps=1e-8, weight_decay=0, momentum=0, centered=False)

RMSprop（Root Mean Square Propagation）优化器

自适应学习率优化器，特别适用于递归神经网络（RNN）。它通过维护梯度平方的移动平均值来调整每个参数的学习率。

参数:

params (Iterable[riemann.TN or riemann.nn.Parameter] or List[Dict[str, Any]]) – 待优化参数的迭代器或定义参数组的字典列表
lr (float, optional) – 学习率
alpha (float, optional) – 平滑常数，用于计算梯度平方的指数移动平均值
eps (float, optional) – 添加到分母以提高数值稳定性的项
weight_decay (float, optional) – 权重衰减（L2正则化）系数
momentum (float, optional) – 动量因子
centered (bool, optional) – 是否使用中心化的RMSprop（使用梯度的移动平均值）

step(): 执行单个优化步骤

学习率调度器

class riemann.optim.lr_scheduler.LRScheduler(optimizer, last_epoch=-1, verbose=False)

所有学习率调度器的基类

参数:

optimizer (riemann.optim.Optimizer) – 要调整学习率的优化器
last_epoch (int, optional) – 最后一个epoch的索引
verbose (bool, optional) – 是否打印学习率更新

step(epoch=None)

执行单个调度步骤

参数:: epoch (int, optional) – 当前epoch索引

get_lr()

返回当前epoch的学习率

返回:: 每个参数组的学习率
返回类型:: List[float]

get_last_lr()

返回最后计算的学习率

返回:: 每个参数组的学习率
返回类型:: List[float]

state_dict()

返回调度器的状态字典

返回:: 调度器状态
返回类型:: Dict[str, Any]

load_state_dict(state_dict)

加载调度器状态

参数:: state_dict (Dict[str, Any]) – 要加载的状态字典

class riemann.optim.lr_scheduler.StepLR(optimizer, step_size, gamma=0.1, last_epoch=-1, verbose=False)

每隔step_size个epoch将学习率衰减一个因子

参数:

optimizer (riemann.optim.Optimizer) – 要调整学习率的优化器
step_size (int) – 学习率衰减的周期
gamma (float, optional) – 学习率衰减的乘法因子
last_epoch (int, optional) – 最后一个epoch的索引
verbose (bool, optional) – 是否打印学习率更新

class riemann.optim.lr_scheduler.MultiStepLR(optimizer, milestones, gamma=0.1, last_epoch=-1, verbose=False)

在指定的milestones处将学习率衰减一个因子

参数:

optimizer (riemann.optim.Optimizer) – 要调整学习率的优化器
milestones (List[int]) – epoch索引列表
gamma (float, optional) – 学习率衰减的乘法因子
last_epoch (int, optional) – 最后一个epoch的索引
verbose (bool, optional) – 是否打印学习率更新

class riemann.optim.lr_scheduler.ExponentialLR(optimizer, gamma, last_epoch=-1, verbose=False)

指数衰减学习率

参数:

optimizer (riemann.optim.Optimizer) – 要调整学习率的优化器
gamma (float) – 学习率衰减的乘法因子
last_epoch (int, optional) – 最后一个epoch的索引
verbose (bool, optional) – 是否打印学习率更新

class riemann.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max, eta_min=0, last_epoch=-1, verbose=False)

使用余弦函数退火学习率

参数:

optimizer (riemann.optim.Optimizer) – 要调整学习率的优化器
T_max (int) – 最大迭代次数
eta_min (float, optional) – 最小学习率
last_epoch (int, optional) – 最后一个epoch的索引
verbose (bool, optional) – 是否打印学习率更新

class riemann.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10, verbose=False, threshold=1e-4, threshold_mode='rel', cooldown=0, min_lr=0, eps=1e-8)

当指标停止改善时减少学习率

参数:

optimizer (riemann.optim.Optimizer) – 要调整学习率的优化器
mode (str, optional) – ‘min’或’max’之一
factor (float, optional) – 学习率衰减的乘法因子
patience (int, optional) – 学习率将在多少个没有改善的epoch后减少
verbose (bool, optional) – 是否打印学习率更新
threshold (float, optional) – 测量新最优值的阈值
threshold_mode (str, optional) – ‘rel’或’abs’之一
cooldown (int, optional) – 学习率减少后恢复正常操作前要等待的epoch数
min_lr (float or List[float], optional) – 最小学习率
eps (float, optional) – 应用于lr的最小衰减

step(metrics, epoch=None)

执行单个调度步骤

参数:

metrics (float) – 要检查的指标值
epoch (int, optional) – 当前epoch索引