def testLoweringCache(self): ctx = frontend.Context() def my_type(args): return args[0] def function(x): return x + x ctx.batch_uncurried(function, my_type) sig = [instructions.TensorType(np.int64, ())] prog1 = ctx.program_lowered('function', sig, NP_BACKEND) prog2 = ctx.program_lowered('function', sig, NP_BACKEND) self.assertEqual(prog1, prog2) sig2 = [instructions.TensorType(np.int32, ())] prog3 = ctx.program_lowered('function', sig2, NP_BACKEND) self.assertNotEqual(prog1, prog3)
def synthetic_pattern_variable_program(include_types=True): """A program that tests product types. Args: include_types: If False, we omit types on the variables, requiring a type inference pass. Returns: program: `instructions.Program`. """ block = instructions.Block([ instructions.PrimOp(["inp"], "many", lambda x: (x + 1, (x + 2, x + 3))), instructions.PrimOp(["many"], ["one", "two"], lambda x: x), ], instructions.halt_op()) leaf = instructions.TensorType(np.int64, ()) the_vars = { "inp": instructions.Type(leaf), "many": instructions.Type((leaf, (leaf, leaf))), "one": instructions.Type(leaf), "two": instructions.Type((leaf, leaf)), } if not include_types: _strip_types(the_vars) return instructions.Program(instructions.ControlFlowGraph([block]), [], the_vars, ["inp"], "two")
def _merge_tensor_type(old_type, obtained_type, backend): """Merges an updated instructions.TensorType for a single field.""" if old_type is None: return obtained_type # Update the inferred dtype. if obtained_type is None: return old_type # One of old_type or obtained_type must be a TensorType because that's when # the caller's pattern_map2 would call _merge_tensor_type. if not isinstance(old_type, instructions.TensorType): raise ValueError( 'Type mismatch: Expected struct type {}, got {}'.format( old_type, obtained_type)) if not isinstance(obtained_type, instructions.TensorType): raise ValueError( 'Type mismatch: Expected tensor type {}, got {}'.format( old_type, obtained_type)) dtype = old_type.dtype obtained_dtype = obtained_type.dtype if dtype is None: dtype = obtained_dtype elif obtained_dtype is not None: dtype = backend.merge_dtypes(dtype, obtained_dtype) # Update the inferred shape. shape = old_type.shape obtained_shape = obtained_type.shape if shape is None: shape = obtained_shape elif obtained_shape is not None: shape = backend.merge_shapes(shape, obtained_shape) return instructions.TensorType(dtype, shape)
def add_batch_dim_one_var(type_): return instructions.Type( instructions.pattern_map( lambda t: instructions.TensorType(t.dtype, (new_batch_dim, ) + t.shape), type_.tensors, leaf_type=instructions.TensorType))
def testFibonacciTF(self): batch_fibo = frontend.Context().batch_uncurried( fibonacci, lambda *args: instructions.TensorType(np.int64, ())) input_2 = self._build_tensor(np.array([6, 7, 8], dtype=np.int64)) answer = batch_fibo(input_2, max_stack_depth=15, backend=TF_BACKEND) self._check_batch_size(answer, 3) self.assertAllEqual([8, 13, 21], self.evaluate(answer))
def var_init(max_stack_depth, initial_value): type_ = inst.TensorType(initial_value.dtype, initial_value.shape[1:]) var = TF_BACKEND.create_variable( None, inst.VariableAllocation.FULL, type_, max_stack_depth, batch_size=initial_value.shape[0]) return var.update( initial_value, TF_BACKEND.full_mask(initial_value.shape[0]))
def testSyncingAsExpected(self): execution_counter_box = [0] def count_executions(x): execution_counter_box[0] += 1 return x # Instrumented test program def fibonacci_inst(n): if n <= 1: return 1 else: left = fibonacci_inst(n - 2) right = fibonacci_inst(n - 1) return count_executions(left + right) batch_fibo = frontend.Context().batch_uncurried( fibonacci_inst, lambda *args: instructions.TensorType(np.int64, ())) self.assertEqual([3, 21, 5, 8], list( batch_fibo(np.array([3, 7, 4, 5], dtype=np.int64), max_stack_depth=15, backend=NP_BACKEND))) # Expect 22 executions # - 2 for type checking, plus # - 20 because that's how many times 1 needs to be added to # 1 to get 21. self.assertEqual(22, execution_counter_box[0])
def testFibonacciNumpy(self): batch_fibo = frontend.Context().batch_uncurried( fibonacci, lambda *args: instructions.TensorType(np.int64, ())) self.assertEqual( [13, 21, 34, 55], list(batch_fibo(np.array([6, 7, 8, 9], dtype=np.int64), max_stack_depth=15, backend=NP_BACKEND)))
def testFibonacciNumpyStackless(self): if not tf.executing_eagerly(): return batch_fibo = frontend.Context().batch_uncurried( fibonacci, lambda *args: instructions.TensorType(np.int64, ())) self.assertEqual( [3, 21, 5, 8], list(batch_fibo(np.array([3, 7, 4, 5], dtype=np.int64), max_stack_depth=15, backend=NP_BACKEND, stackless=True)))
def type_of(self, t, preferred_dtype=None): """Returns the `instructions.Type` of `t`. Args: t: `np.ndarray` or a Python constant. preferred_dtype: dtype to prefer, if `t` is a constant. Returns: vm_type: `instructions.TensorType` describing `t` """ t = np.array(t, dtype=preferred_dtype) return instructions.TensorType(t.dtype, t.shape)
def type_of(self, t, dtype_hint=None): """Returns the `instructions.Type` of `t`. Args: t: `tf.Tensor` or a Python or numpy constant. dtype_hint: dtype to prefer, if `t` is a constant. Returns: vm_type: `instructions.TensorType` describing `t`. """ if tf.executing_eagerly(): new_t = tf.convert_to_tensor(value=t, dtype=dtype_hint) else: with tf.Graph().as_default(): # Use a scratch graph. new_t = tf.convert_to_tensor(value=t, dtype=dtype_hint) dtype = new_t.dtype.base_dtype.as_numpy_dtype shape = None if new_t.shape.ndims is None else tuple(new_t.shape.as_list()) return instructions.TensorType(dtype, shape)
def my_type(args): return instructions.TensorType(args[0].dtype, ())
def _strip_batch_dim(type_): return instructions.pattern_map( lambda t: instructions.TensorType(t.dtype, t.shape[1:]), type_, leaf_type=instructions.TensorType)
def pred_type(_): return instructions.TensorType(np.bool, ())
def pred_type(t): return instructions.TensorType(np.bool, t[0].shape)
def testClosingOverTensorDoesntRaise(self): x = tf.constant(0.) def f(y): return y * x arg_types = [inst.Type([inst.TensorType(shape=[], dtype=np.float32)])] TF_BACKEND.run_on_dummies(f, arg_types)
def pred_type(_): return instructions.TensorType(np.int32, ())
def my_type(_): int_ = instructions.TensorType(np.int64, ()) return ((int_, int_), int_, (int_, int_))