jax.experimental.sparse.bcoo_更新布局

jax.experimental.sparse.bcoo_更新布局#

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) – optional(int) 输出矩阵中的批处理维度数量。如果为 None，则 n_batch = mat.n_batch。

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

  • on_inefficient (str | None) – optional(string), 可选值为 ['error', 'warn', None]。指定在发生低效重新配置时的行为。低效重新配置定义为结果表示的大小远大于输入表示的大小时的重新配置。

返回:

BCOO 数组 一个 BCOO 数组，表示与输入相同的稀疏数组，具有指定的布局。mat_out.todense() 将与 mat.todense() 在适当的精度内匹配。

返回类型:

mat_out

目录