jax.experimental.sparse.bcoo_update_layout

jax.experimental.sparse.bcoo_update_layout#

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

更新 BCOO 矩阵的存储布局(即 n_batch 和 n_dense)。

在许多情况下,这可以在不引入过多存储开销的情况下完成。但是,增加 mat.n_batchmat.n_dense 将导致非常低效的存储,具有许多显式存储的零,除非新的批次或密集维度的大小为 0 或 1。在这种情况下,bcoo_update_layout 将引发 SparseEfficiencyError。可以通过指定 on_inefficient 参数来抑制此警告。

参数:
  • mat (BCOO) – BCOO 数组

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

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

  • on_inefficient (str | None) – 可选(字符串),取值为 ['error', 'warn', None] 中的一个。指定在进行低效重构时的行为。低效重构指的是,重构后的表示的大小远大于输入表示的大小。

返回:

BCOO 数组

一个 BCOO 数组,表示与输入相同的稀疏数组,但具有指定的布局。 mat_out.todense() 将与 mat.todense() 相匹配,直至达到相应的精度。

返回类型:

mat_out