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 | None) – 可选(int) 输出矩阵中的批处理维度数。如果为 None,则 n_batch = mat.n_batch。
n_dense (int | None | None) – 可选(int) 输出矩阵中的密集维度数。如果为 None,则 n_dense = mat.n_dense。
on_inefficient (str | None) – 可选(string),
['error', 'warn', None]
之一。指定在低效重新配置情况下的行为。这被定义为一种重新配置,其中结果表示的大小远大于输入表示的大小。
- 返回:
- BCOO 数组
一个 BCOO 数组,表示与输入相同的稀疏数组,并具有指定的布局。
mat_out.todense()
将匹配mat.todense()
,精度适当。
- 返回类型:
mat_out