def _dot_general_sparse(spenv, *argspecs, dimension_numbers, precision,
preferred_element_type):
A, B = argspecs_to_arrays(spenv, argspecs)
if argspecs[0].is_sparse() and argspecs[1].is_sparse():
shape = sparse.bcoo._dot_general_validated_shape(
A.shape, B.shape, dimension_numbers)
data, indices = sparse.bcoo_spdot_general(
A.data,
A.indices,
B.data,
B.indices,
lhs_shape=A.shape,
rhs_shape=B.shape,
dimension_numbers=dimension_numbers)
return [ArgSpec(shape, spenv.push(data), spenv.push(indices))]
elif argspecs[0].is_sparse():
result = sparse.bcoo_dot_general(A.data,
A.indices,
B,
lhs_shape=A.shape,
dimension_numbers=dimension_numbers)
else:
result = sparse.bcoo_rdot_general(A,
B.data,
B.indices,
rhs_shape=B.shape,
dimension_numbers=dimension_numbers)
return [ArgSpec(result.shape, spenv.push(result), None)]
def _dot_general_sparse(spenv, *spvalues, dimension_numbers, precision,
preferred_element_type):
# TODO(jakevdp): pass along these unused configurations?
del precision, preferred_element_type # unused
result = sparse.bcoo_dot_general(*spvalues_to_arrays(spenv, spvalues),
dimension_numbers=dimension_numbers)
return arrays_to_spvalues(spenv, [result])
def _dot_general_sparse(spenv, *spvalues, dimension_numbers, precision, preferred_element_type):
A, B = spvalues_to_arrays(spenv, spvalues)
if spvalues[0].is_sparse() and spvalues[1].is_sparse():
shape = sparse.bcoo._dot_general_validated_shape(A.shape, B.shape, dimension_numbers)
data, indices = sparse.bcoo_spdot_general(A, B, dimension_numbers=dimension_numbers)
return [spenv.sparse(shape, data, indices)]
elif spvalues[0].is_sparse():
result = sparse.bcoo_dot_general(A, B, dimension_numbers=dimension_numbers)
else:
result = sparse.bcoo_rdot_general(A, B, dimension_numbers=dimension_numbers)
return [spenv.dense(result)]
def _dot_general_sparse(spenv, *argspecs, dimension_numbers, precision, preferred_element_type):
if argspecs[0].is_sparse() and argspecs[1].is_sparse():
raise NotImplementedError("dot_general between two sparse matrices.")
A, B = argspecs_to_arrays(spenv, argspecs)
if argspecs[0].is_sparse():
result = sparse.bcoo_dot_general(A.data, A.indices, B, lhs_shape=A.shape,
dimension_numbers=dimension_numbers)
else:
result = sparse.bcoo_rdot_general(A, B.data, B.indices, rhs_shape=B.shape,
dimension_numbers=dimension_numbers)
return [ArgSpec(result.shape, spenv.push(result), None)]
def _dot_general_sparse(spenv, *argspecs, dimension_numbers, precision,
preferred_element_type):
assert len(argspecs) == 2
assert argspecs[0].is_sparse() and not argspecs[1].is_sparse()
args = argspecs_to_arrays(spenv, argspecs)
result = sparse.bcoo_dot_general(args[0].data,
args[0].indices,
args[1],
lhs_shape=args[0].shape,
dimension_numbers=dimension_numbers)
argspec = ArgSpec(result.shape, spenv.push(result), None)
return [argspec]
def f_sparse(Y):
return sparse.bcoo_dot_general(data, indices, Y, lhs_shape=X.shape,
dimension_numbers=dimension_numbers)
def f_sparse(data, indices, lhs, rhs):
return sparse.bcoo_dot_general(data, indices, rhs,
lhs_shape=lhs.shape,
dimension_numbers=dimension_numbers)