def test_infer_shape(self):
rng = np.random.RandomState(3453)
adtens4 = dtensor4()
aivec = ivector()
aivec_val = [3, 4, 2, 5]
adtens4_val = rng.rand(*aivec_val)
self._compile_and_check(
[adtens4, aivec],
[SpecifyShape()(adtens4, aivec)],
[adtens4_val, aivec_val],
SpecifyShape,
)
def test_rebroadcast(self):
rng = np.random.RandomState(3453)
# Rebroadcast
adtens4 = dtensor4()
adict = [(0, False), (1, True), (2, False), (3, True)]
adtens4_val = rng.rand(2, 1, 3, 1).astype(config.floatX)
self._compile_and_check(
[adtens4],
[Rebroadcast(*adict)(adtens4)],
[adtens4_val],
Rebroadcast,
warn=False,
)
adtens4_bro = TensorType("float64", (True, True, True, False))()
bdict = [(0, True), (1, False), (2, False), (3, False)]
adtens4_bro_val = rng.rand(1, 1, 1, 3).astype(config.floatX)
self._compile_and_check(
[adtens4_bro],
[Rebroadcast(*bdict)(adtens4_bro)],
[adtens4_bro_val],
Rebroadcast,
)