def test_advanced_indexing_lazy(output_shape): x = Tensor(randn((2, 3, 4) + output_shape), OrderedDict([ ('i', bint(2)), ('j', bint(3)), ('k', bint(4)), ])) u = Variable('u', bint(2)) v = Variable('v', bint(3)) with interpretation(lazy): i = Number(1, 2) - u j = Number(2, 3) - v k = u + v expected_data = empty((2, 3) + output_shape) i_data = x.materialize(i).data j_data = x.materialize(j).data k_data = x.materialize(k).data for u in range(2): for v in range(3): expected_data[u, v] = x.data[i_data[u], j_data[v], k_data[u, v]] expected = Tensor(expected_data, OrderedDict([ ('u', bint(2)), ('v', bint(3)), ])) assert_equiv(expected, x(i, j, k)) assert_equiv(expected, x(i=i, j=j, k=k)) assert_equiv(expected, x(i=i, j=j)(k=k)) assert_equiv(expected, x(j=j, k=k)(i=i)) assert_equiv(expected, x(k=k, i=i)(j=j)) assert_equiv(expected, x(i=i)(j=j, k=k)) assert_equiv(expected, x(j=j)(k=k, i=i)) assert_equiv(expected, x(k=k)(i=i, j=j)) assert_equiv(expected, x(i=i)(j=j)(k=k)) assert_equiv(expected, x(i=i)(k=k)(j=j)) assert_equiv(expected, x(j=j)(i=i)(k=k)) assert_equiv(expected, x(j=j)(k=k)(i=i)) assert_equiv(expected, x(k=k)(i=i)(j=j)) assert_equiv(expected, x(k=k)(j=j)(i=i))
def diff_fn(p_data): p = Tensor(p_data, be_inputs) q = p.sample(sampled_vars, sample_inputs, rng_key=rng_key) mq = p.materialize(q).reduce(ops.logaddexp, 'n') mq = mq.align(tuple(p.inputs)) _, (p_data, mq_data) = align_tensors(p, mq) assert p_data.shape == mq_data.shape return (ops.exp(mq_data) * probe).sum() - (ops.exp(p_data) * probe).sum(), mq
def test_diagonal_rename(): x = Tensor( randn(2, 2, 3), OrderedDict(a=funsor.Bint[2], b=funsor.Bint[2], c=funsor.Bint[3]), 'real') d = Variable("d", funsor.Bint[2]) dt = x.materialize(d) yt = x(a=dt, b=dt) y = x(a=d, b=d) assert_close(y, yt)
def eager_subs(self, subs): assert isinstance(subs, tuple) prototype = Tensor(self.info_vec) subs = tuple( (k, v if isinstance(v, (Variable, Slice)) else prototype.materialize(v)) for k, v in subs if k in self.inputs) if not subs: return self # Constants and Affine funsors are eagerly substituted; # everything else is lazily substituted. lazy_subs = tuple( (k, v) for k, v in subs if not isinstance(v, (Number, Tensor, Variable, Slice)) and not (is_affine(v) and affine_inputs(v))) var_subs = tuple((k, v) for k, v in subs if isinstance(v, Variable)) int_subs = tuple((k, v) for k, v in subs if isinstance(v, (Number, Tensor, Slice)) if v.dtype != 'real') real_subs = tuple((k, v) for k, v in subs if isinstance(v, (Number, Tensor)) if v.dtype == 'real') affine_subs = tuple((k, v) for k, v in subs if is_affine(v) and affine_inputs(v) and not isinstance(v, Variable)) if var_subs: return self._eager_subs_var( var_subs, int_subs + real_subs + affine_subs + lazy_subs) if int_subs: return self._eager_subs_int(int_subs, real_subs + affine_subs + lazy_subs) if real_subs: return self._eager_subs_real(real_subs, affine_subs + lazy_subs) if affine_subs: return self._eager_subs_affine(affine_subs, lazy_subs) return reflect(Subs, self, lazy_subs)