def test_magma_opt_float16(self):
ops_to_gpu = [
(MatrixInverse(), GpuMagmaMatrixInverse),
(SVD(), GpuMagmaSVD),
(QRFull(mode="reduced"), GpuMagmaQR),
(QRIncomplete(mode="r"), GpuMagmaQR),
# TODO: add support for float16 to Eigh numpy
# (Eigh(), GpuMagmaEigh),
(Cholesky(), GpuMagmaCholesky),
]
for op, gpu_op in ops_to_gpu:
A = theano.tensor.matrix("A", dtype="float16")
fn = theano.function([A], op(A), mode=mode_with_gpu.excluding("cusolver"))
assert any(
[isinstance(node.op, gpu_op) for node in fn.maker.fgraph.toposort()]
)
from theano.tensor.nlinalg import QRFull
from unification import var
from kanren import eq
from kanren.core import lall
from kanren.graph import applyo
from kanren.constraints import neq
from etuples import etuple, etuplize
from ...theano.meta import mt
mt.nlinalg.qr_full = mt(QRFull("reduced"))
owner_inputs = attrgetter("owner.inputs")
normal_get_size = toolz.compose(itemgetter(2), owner_inputs)
normal_get_rng = toolz.compose(itemgetter(3), owner_inputs)
def update_name_suffix(x, old_x, suffix): # pragma: no cover
new_name = old_x.name + suffix
x.name = new_name
return x
def normal_normal_regression(Y, X, beta, Y_args_tail=None, beta_args=None):
"""Create a goal for a normal-normal regression of the form `Y ~ N(X * beta, sd**2)`."""
Y_args_tail = Y_args_tail or var()
beta_args = beta_args or var()