jax.experimental.sparse.bcoo_update_layout#

jax.experimental.sparse.bcoo_update_layout(mat, *, n_batch=None, n_dense=None, on_inefficient='error')[源代码]#

更新 BCOO 矩阵的存储布局（即 n_batch 和 n_dense）。

在许多情况下，这可以在不引入不当存储开销的情况下完成。然而，增加 mat.n_batch 或 mat.n_dense 将导致非常低效的存储，其中包含许多显式存储的零，除非新的批次或密集维度的尺寸为 0 或 1。在这种情况下，bcoo_update_layout 将引发 SparseEfficiencyError。这可以通过指定 on_inefficient 参数来抑制。

参数:

  • mat (BCOO) – BCOO 数组

  • n_batch (int | None) – 可选(int) 输出矩阵中的批次维度数量。如果为 None，则 n_batch = mat.n_batch。

  • n_dense (int | None) – 可选(int) 输出矩阵中的密集维度数量。如果为 None，则 n_dense = mat.n_dense。

  • on_inefficient (str | None) – 可选(字符串)，以下选项之一：['error', 'warn', None]。指定在配置效率低下时应采取的行为。这种情况被定义为一种配置，其中结果表示的大小远大于输入表示的大小。

返回:

BCOO 数组

一个 BCOO 数组，表示与输入相同的稀疏数组，并具有指定的布局。mat_out.todense() 将匹配 mat.todense()，但精度可能略有不同。

返回类型:

mat_out