コード例 #1
0
def _broadcast_in_dim_sparse(spenv, *spvalues, shape, broadcast_dimensions):
    operand, = spvalues
    operand_promoted = spvalues_to_arrays(spenv, operand)
    mat = sparse.bcoo_broadcast_in_dim(
        operand_promoted,
        shape=shape,
        broadcast_dimensions=broadcast_dimensions)
    return (spenv.sparse(shape, mat.data, mat.indices), )
コード例 #2
0
ファイル: transform.py プロジェクト: ahoenselaar/jax 
def _broadcast_in_dim_sparse(spenv, *argspecs, shape, broadcast_dimensions):
    operand, = argspecs
    data, indices = sparse.bcoo_broadcast_in_dim(
        operand.data(spenv),
        operand.indices(spenv),
        shape=operand.shape,
        new_shape=shape,
        broadcast_dimensions=broadcast_dimensions)
    return (ArgSpec(shape, spenv.push(data), spenv.push(indices)), )