def _broadcast( cls, shape1: TensorFluentShape, shape2: TensorFluentShape ) -> TensorFluentShape: s1, s2 = TensorFluentShape.broadcast(shape1, shape2) s1 = s1 if s1 is not None else shape1.as_list() s2 = s2 if s2 is not None else shape2.as_list() x1, x2 = np.zeros(s1), np.zeros(s2) y = np.broadcast(x1, x2) return TensorFluentShape(y.shape, batch=(shape1.batch or shape2.batch))
def __init__(self, tensor: tf.Tensor, scope: List[str], batch: bool = False) -> None: self.tensor = tensor self.scope = TensorFluentScope(scope) self.shape = TensorFluentShape(tensor.shape, batch)
def test_batch_normal(compiler): with compiler.graph.as_default(): shape_scope = { "mu/1": TensorFluentShape((64, 16), batch=True), "sigma/1": TensorFluentShape((64, 16), batch=True), } scope = { "mu/1": TensorFluent(tf.zeros([64, 16]), scope=["?x"], batch=True), "sigma/1": TensorFluent(tf.ones([64, 16]), scope=["?x"], batch=True), } noise1 = compiler._get_expression_reparameterization(X1, scope=shape_scope) _test_reparameterization_dist(noise1, [(NORMAL, [64, 16])]) _test_reparameterized_expression(compiler, X1, scope=scope, noise=noise1, name="noise1") noise2 = compiler._get_expression_reparameterization(X2, scope=shape_scope) _test_reparameterization_dist(noise2, [(NORMAL, [64, 16])]) _test_reparameterized_expression(compiler, X2, scope=scope, noise=noise2, name="noise2") noise3 = compiler._get_expression_reparameterization(X3, scope=shape_scope) _test_reparameterization_dist(noise3, [(NORMAL, [64, 16])]) _test_reparameterized_expression(compiler, X3, scope=scope, noise=noise3, name="noise3")
def test_exponential(compiler): # rainfall(?r) = Exponential(RAIN_RATE(?r)); with compiler.graph.as_default(): shape_scope = {"rate/1": TensorFluentShape((32, 8), batch=True)} scope = { "rate/1": TensorFluent(tf.ones((32, 8)), scope=["?r"], batch=True) } noise1 = compiler._get_expression_reparameterization(EXP1, scope=shape_scope) _test_reparameterization_dist(noise1, [(UNIFORM, [32, 8])]) _test_reparameterized_expression(compiler, EXP1, scope=scope, noise=noise1, name="noise1")
def _binary_op(cls, x: 'TensorFluent', y: 'TensorFluent', op: Callable[[tf.Tensor, tf.Tensor], tf.Tensor], dtype: tf.DType) -> 'TensorFluent': '''Returns a TensorFluent for the binary `op` applied to fluents `x` and `y`. Args: x: The first operand. y: The second operand. op: The binary operator. dtype: The output's data type. Returns: A TensorFluent wrapping the binary operator's output. ''' # scope s1 = x.scope.as_list() s2 = y.scope.as_list() scope, perm1, perm2 = TensorFluentScope.broadcast(s1, s2) if x.batch and perm1 != []: perm1 = [0] + [p + 1 for p in perm1] if y.batch and perm2 != []: perm2 = [0] + [p + 1 for p in perm2] x = x.transpose(perm1) y = y.transpose(perm2) # shape reshape1, reshape2 = TensorFluentShape.broadcast(x.shape, y.shape) if reshape1 is not None: x = x.reshape(reshape1) if reshape2 is not None: y = y.reshape(reshape2) # dtype x = x.cast(dtype) y = y.cast(dtype) # operation t = op(x.tensor, y.tensor) # batch batch = x.batch or y.batch return TensorFluent(t, scope, batch=batch)
def test_function(compiler): with compiler.graph.as_default(): shape_scope = { "mu/1": TensorFluentShape([24], batch=False), "sigma/1": TensorFluentShape([24], batch=False), } scope = { "mu/1": TensorFluent(tf.zeros([24]), scope=["?x"], batch=False), "sigma/1": TensorFluent(tf.ones([24]), scope=["?x"], batch=False), } noise1 = compiler._get_expression_reparameterization(EXP_2, scope=shape_scope) _test_reparameterization_dist(noise1, []) _test_reparameterized_expression(compiler, EXP_2, scope=scope, noise=noise1, name="noise1") noise2 = compiler._get_expression_reparameterization(EXP_Z, scope=shape_scope) _test_reparameterization_dist(noise2, [(NORMAL, [1])]) _test_reparameterized_expression(compiler, EXP_Z, scope=scope, noise=noise2, name="noise2") noise3 = compiler._get_expression_reparameterization(EXP_X1, scope=shape_scope) _test_reparameterization_dist(noise3, [(NORMAL, [24])]) _test_reparameterized_expression(compiler, EXP_X1, scope=scope, noise=noise3, name="noise3") noise4 = compiler._get_expression_reparameterization(Y1, scope=shape_scope) _test_reparameterization_dist(noise4, [(NORMAL, [1]), (NORMAL, [1])]) _test_reparameterized_expression(compiler, Y1, scope=scope, noise=noise4, name="noise4") noise5 = compiler._get_expression_reparameterization(Y2, scope=shape_scope) _test_reparameterization_dist(noise5, [(NORMAL, [1]), (NORMAL, [24])]) _test_reparameterized_expression(compiler, Y2, scope=scope, noise=noise5, name="noise5") noise6 = compiler._get_expression_reparameterization(Y3, scope=shape_scope) _test_reparameterization_dist(noise6, [(NORMAL, [24]), (NORMAL, [24])]) _test_reparameterized_expression(compiler, Y3, scope=scope, noise=noise6, name="noise6")
def test_arithmetic(compiler): with compiler.graph.as_default(): shape_scope = { "mu/1": TensorFluentShape([32], batch=False), "sigma/1": TensorFluentShape([32], batch=False), } scope = { "mu/1": TensorFluent(tf.zeros([32]), scope=["?x"], batch=False), "sigma/1": TensorFluent(tf.ones([32]), scope=["?x"], batch=False), } noise1 = compiler._get_expression_reparameterization(TWO, scope={}) _test_reparameterization_dist(noise1, []) _test_reparameterized_expression(compiler, TWO, scope={}, noise=noise1, name="noise1") noise2 = compiler._get_expression_reparameterization(Z_TIMES_Z, scope={}) _test_reparameterization_dist(noise2, [(NORMAL, [1]), (NORMAL, [1])]) _test_reparameterized_expression(compiler, Z_TIMES_Z, scope={}, noise=noise2, name="noise2") noise3 = compiler._get_expression_reparameterization(X2_TIMES_X2, scope=shape_scope) _test_reparameterization_dist(noise3, [(NORMAL, [32]), (NORMAL, [32])]) _test_reparameterized_expression(compiler, X2_TIMES_X2, scope=scope, noise=noise3, name="noise3") noise4 = compiler._get_expression_reparameterization(MU_PLUS_Z, scope=shape_scope) _test_reparameterization_dist(noise4, [(NORMAL, [1])]) _test_reparameterized_expression(compiler, MU_PLUS_Z, scope=scope, noise=noise4, name="noise4") noise5 = compiler._get_expression_reparameterization(Z_PLUS_MU, scope=shape_scope) _test_reparameterization_dist(noise5, [(NORMAL, [1])]) _test_reparameterized_expression(compiler, Z_PLUS_MU, scope=scope, noise=noise5, name="noise5") noise6 = compiler._get_expression_reparameterization(MU_PLUS_X2, scope=shape_scope) _test_reparameterization_dist(noise6, [(NORMAL, [32])]) _test_reparameterized_expression(compiler, MU_PLUS_X2, scope=scope, noise=noise6, name="noise6") noise7 = compiler._get_expression_reparameterization(X2_PLUS_MU, scope=shape_scope) _test_reparameterization_dist(noise7, [(NORMAL, [32])]) _test_reparameterized_expression(compiler, X2_PLUS_MU, scope=scope, noise=noise7, name="noise7") noise8 = compiler._get_expression_reparameterization(X1_PLUS_Z, scope=shape_scope) _test_reparameterization_dist(noise8, [(NORMAL, [32]), (NORMAL, [1])]) _test_reparameterized_expression(compiler, X1_PLUS_Z, scope=scope, noise=noise8, name="noise8") noise9 = compiler._get_expression_reparameterization(Z_PLUS_X1, scope=shape_scope) _test_reparameterization_dist(noise9, [(NORMAL, [1]), (NORMAL, [32])]) _test_reparameterized_expression(compiler, Z_PLUS_X1, scope=scope, noise=noise9, name="noise9")
def test_multivariate_normal(compiler): with compiler.graph.as_default(): shape_scope = { "mu/1": TensorFluentShape([32], batch=False), "sigma/1": TensorFluentShape([32], batch=False), } scope = { "mu/1": TensorFluent(tf.zeros([32]), scope=["?x"], batch=False), "sigma/1": TensorFluent(tf.ones([32]), scope=["?x"], batch=False), } noise1 = compiler._get_expression_reparameterization(X1, scope=shape_scope) _test_reparameterization_dist(noise1, [(NORMAL, [32])]) _test_reparameterized_expression(compiler, X1, scope=scope, noise=noise1, name="noise1") noise2 = compiler._get_expression_reparameterization(X2, scope=shape_scope) _test_reparameterization_dist(noise2, [(NORMAL, [32])]) _test_reparameterized_expression(compiler, X2, scope=scope, noise=noise2, name="noise2") noise3 = compiler._get_expression_reparameterization(X3, scope=shape_scope) _test_reparameterization_dist(noise3, [(NORMAL, [32])]) _test_reparameterized_expression(compiler, X3, scope=scope, noise=noise3, name="noise3") noise4 = compiler._get_expression_reparameterization(X4, scope=shape_scope) _test_reparameterization_dist(noise4, [(NORMAL, [1]), (NORMAL, [1])]) _test_reparameterized_expression(compiler, X4, scope=scope, noise=noise4, name="noise4") noise5 = compiler._get_expression_reparameterization(X5, scope=shape_scope) _test_reparameterization_dist(noise5, [(NORMAL, [32]), (NORMAL, [32])]) _test_reparameterized_expression(compiler, X5, scope=scope, noise=noise5, name="noise5") noise6 = compiler._get_expression_reparameterization(X6, scope=shape_scope) _test_reparameterization_dist(noise6, [(NORMAL, [32]), (NORMAL, [32])]) _test_reparameterized_expression(compiler, X6, scope=scope, noise=noise6, name="noise6") noise7 = compiler._get_expression_reparameterization(X7, scope=shape_scope) _test_reparameterization_dist(noise7, [(NORMAL, [1]), (NORMAL, [1]), (NORMAL, [1])]) _test_reparameterized_expression(compiler, X7, scope=scope, noise=noise7, name="noise7") noise8 = compiler._get_expression_reparameterization(X8, scope=shape_scope) _test_reparameterization_dist(noise8, [(NORMAL, [1]), (NORMAL, [1]), (NORMAL, [32])]) _test_reparameterized_expression(compiler, X8, scope=scope, noise=noise8, name="noise8") noise9 = compiler._get_expression_reparameterization(X9, scope=shape_scope) _test_reparameterization_dist(noise9, [(NORMAL, [32]), (NORMAL, [1]), (NORMAL, [1]), (NORMAL, [32])]) _test_reparameterized_expression(compiler, X9, scope=scope, noise=noise9, name="noise9")
def _get_reparameterization( self, expr: Expression, scope: ShapeScope, noise: NoiseList ) -> TensorFluentShape: etype = expr.etype args = expr.args if etype[0] == "constant": return TensorFluentShape([1], batch=False) elif etype[0] == "pvar": name = expr._pvar_to_name(args) if name not in scope: raise ValueError("Variable {} not in scope.".format(name)) shape = scope[name] return shape elif etype[0] == "randomvar": if etype[1] == "Normal": mean_shape = self._get_reparameterization(args[0], scope, noise) var_shape = self._get_reparameterization(args[1], scope, noise) shape = ReparameterizationCompiler._broadcast(mean_shape, var_shape) dist = tf.distributions.Normal(loc=0.0, scale=1.0) noise.append((dist, shape.as_list())) return shape elif etype[1] == "Exponential": rate_shape = self._get_reparameterization(args[0], scope, noise) dist = tf.distributions.Uniform(low=0.0, high=1.0) noise.append((dist, rate_shape.as_list())) return rate_shape elif etype[1] == "Gamma": for fluent in self.rddl.get_dependencies(expr): if fluent.is_state_fluent() or fluent.is_action_fluent(): raise ValueError( f"Expression is not an exogenous event: {expr}" ) shape = [] with self.graph.as_default(): scope = self._scope.non_fluents(self.non_fluents) shape_fluent = self._compile_expression(args[0], scope, noise=None) scale_fluent = self._compile_expression(args[1], scope, noise=None) concentration = shape_fluent.tensor rate = 1 / scale_fluent.tensor dist = tf.distributions.Gamma(concentration, rate) noise.append((dist, shape)) return shape elif etype[1] == "Uniform": low_shape = self._get_reparameterization(args[0], scope, noise) high_shape = self._get_reparameterization(args[1], scope, noise) shape = ReparameterizationCompiler._broadcast(low_shape, high_shape) dist = tf.distributions.Uniform(low=0.0, high=1.0) noise.append((dist, shape.as_list())) return shape elif etype[0] in ["arithmetic", "boolean", "relational"]: op1_shape = self._get_reparameterization(args[0], scope, noise) shape = op1_shape if len(args) > 1: op2_shape = self._get_reparameterization(args[1], scope, noise) shape = ReparameterizationCompiler._broadcast(op1_shape, op2_shape) return shape elif etype[0] == "func": op1_shape = self._get_reparameterization(args[0], scope, noise) shape = op1_shape if len(args) > 1: if len(args) == 2: op2_shape = self._get_reparameterization(args[1], scope, noise) shape = ReparameterizationCompiler._broadcast(op1_shape, op2_shape) else: raise ValueError("Invalid function:\n{}".format(expr)) return shape elif etype[0] == "control": if etype[1] == "if": condition_shape = self._get_reparameterization(args[0], scope, noise) true_case_shape = self._get_reparameterization(args[1], scope, noise) false_case_shape = self._get_reparameterization(args[2], scope, noise) shape = ReparameterizationCompiler._broadcast( condition_shape, true_case_shape ) shape = ReparameterizationCompiler._broadcast(shape, false_case_shape) return shape else: raise ValueError("Invalid control flow expression:\n{}".format(expr)) elif etype[0] == "aggregation": return self._get_reparameterization(args[-1], scope, noise) raise ValueError("Expression type unknown: {}".format(etype))
def _get_reparameterization_shape_scope(self) -> ShapeScope: scope = { name: TensorFluentShape(size, batch=False) for name, (_, size) in self.rddl.fluent_table.items() } return scope
def test_broadcast(self): tests = [ (TensorFluentShape([], False), TensorFluentShape([], False), None, None), (TensorFluentShape([8], False), TensorFluentShape([], False), None, None), (TensorFluentShape([], False), TensorFluentShape([8], False), None, None), (TensorFluentShape([8, 8], False), TensorFluentShape([8], False), None, None), (TensorFluentShape([8], False), TensorFluentShape([8, 8], False), None, None), (TensorFluentShape([100], True), TensorFluentShape([100], True), None, None), (TensorFluentShape([100, 8], True), TensorFluentShape([100], True), None, [100, 1]), (TensorFluentShape([100], True), TensorFluentShape([100, 8], True), [100, 1], None), (TensorFluentShape([100, 8, 8], True), TensorFluentShape([100], True), None, [100, 1, 1]), (TensorFluentShape([100], True), TensorFluentShape([100, 8, 8], True), [100, 1, 1], None), (TensorFluentShape([100, 8, 8], True), TensorFluentShape([100, 8], True), None, [100, 1, 8]), (TensorFluentShape([100, 8], True), TensorFluentShape([100, 8, 8], True), [100, 1, 8], None), (TensorFluentShape([100], True), TensorFluentShape([], False), None, None), (TensorFluentShape([], False), TensorFluentShape([], True), None, None), (TensorFluentShape([100], True), TensorFluentShape([], False), None, None), (TensorFluentShape([100], True), TensorFluentShape([8], False), [100, 1], None), (TensorFluentShape([8], False), TensorFluentShape([100], True), None, [100, 1]), (TensorFluentShape([100], True), TensorFluentShape([8, 7], False), [100, 1, 1], None), (TensorFluentShape([8, 7], False), TensorFluentShape([100], True), None, [100, 1, 1]), (TensorFluentShape([100, 8], True), TensorFluentShape([], False), None, None), (TensorFluentShape([], False), TensorFluentShape([100, 8], True), None, None), (TensorFluentShape([100, 8], True), TensorFluentShape([8], False), None, None), (TensorFluentShape([8], False), TensorFluentShape([100, 8], True), None, None), (TensorFluentShape([100, 8, 7], True), TensorFluentShape([7], False), None, [1, 7]), (TensorFluentShape([7], False), TensorFluentShape([100, 8, 7], True), [1, 7], None), (TensorFluentShape([100, 7, 8], True), TensorFluentShape([7, 8], False), None, None), (TensorFluentShape([7, 8], False), TensorFluentShape([100, 7, 8], True), None, None), (TensorFluentShape([8, 8], False), TensorFluentShape([100, 8], True), None, [100, 1, 8]), (TensorFluentShape([100, 8], True), TensorFluentShape([8, 8], False), [100, 1, 8], None), (TensorFluentShape([2, 2], False), TensorFluentShape([1, 2], True), None, [1, 1, 2]), (TensorFluentShape([1, 2], True), TensorFluentShape([2, 2], False), [1, 1, 2], None), ] for s1, s2, ss1, ss2 in tests: reshape1, reshape2 = TensorFluentShape.broadcast(s1, s2) if ss1 is None: self.assertIsNone(reshape1) else: self.assertListEqual(reshape1, ss1) if ss2 is None: self.assertIsNone(reshape2) else: self.assertListEqual(reshape2, ss2)