def testStringKeyGrad(self): with backprop.GradientTape(persistent=True) as tape: m = map_ops.empty_tensor_map() k = constant_op.constant("key") k2 = constant_op.constant("key2") v = constant_op.constant(2.0) v2 = constant_op.constant(22.0) tape.watch(v2) m = map_ops.tensor_map_insert(m, k2, v2) m = map_ops.tensor_map_insert(m, k, v) s = map_ops.tensor_map_size(m) self.assertAllEqual(s, 2) l = map_ops.tensor_map_lookup(m, k, v.dtype) self.assertAllClose(l, v) m = map_ops.tensor_map_insert(m, k, v2) l2 = map_ops.tensor_map_lookup(m, k, v2.dtype) self.assertAllClose(l2, v2) g = tape.gradient(l2 * 5, v2) self.assertAllEqual(g, 5) m, e = map_ops.tensor_map_erase(m, k, v2.dtype) s = map_ops.tensor_map_size(m) self.assertAllEqual(s, 1) self.assertAllClose(e, v2) g2 = tape.gradient(e * 6, v2) self.assertAllEqual(g2, 6) del tape
def testMultipleInsertLookupGrad(self): with backprop.GradientTape(persistent=True) as tape: m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) k2 = constant_op.constant(2.0) k3 = constant_op.constant(3.0) v = constant_op.constant(11.0) v2 = constant_op.constant(12.0) v3 = constant_op.constant(13.0) tape.watch(v) tape.watch(v2) tape.watch(v3) m = map_ops.tensor_map_insert(m, k, v) m = map_ops.tensor_map_insert(m, k2, v2) m = map_ops.tensor_map_insert(m, k3, v3) l = map_ops.tensor_map_lookup(m, k, v.dtype) l2 = map_ops.tensor_map_lookup(m, k2, v2.dtype) l3 = map_ops.tensor_map_lookup(m, k3, v3.dtype) g = tape.gradient(l * 5, v) g2 = tape.gradient(l2 * 6, v2) g3 = tape.gradient(l3 * 7, v3) self.assertAllEqual(g, 5) self.assertAllEqual(g2, 6) self.assertAllEqual(g3, 7) del tape
def testTensorMapEraseFromEmptyMapFails(self): m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) with self.assertRaisesRegex(errors.InvalidArgumentError, "Trying to erase non-existent item."): m = map_ops.tensor_map_erase(m, k, dtypes.float32) self.evaluate(m)
def testStringKeyGrad(self): with backprop.GradientTape(persistent=True) as tape: m = map_ops.empty_tensor_map() k = constant_op.constant("key") k2 = constant_op.constant("key2") v = constant_op.constant(2.0) v2 = constant_op.constant(22.0) tape.watch(v) tape.watch(v2) m = map_ops.tensor_map_insert(m, k, v) m = map_ops.tensor_map_insert(m, k2, v2) s = map_ops.tensor_map_size(m) self.assertAllEqual(s, 2) # Test lookup and gradient. l = map_ops.tensor_map_lookup(m, k, v.dtype) self.assertAllClose(l, v) self.assertAllClose(tape.gradient(l * 5, v), 5) # Test replace and gradient. m = map_ops.tensor_map_insert(m, k, v2) l2 = map_ops.tensor_map_lookup(m, k, v2.dtype) self.assertAllClose(l2, v2) g = tape.gradient(l2 * 6, v2) self.assertAllEqual(g, 6) # Test erase, has key, and gradient. m = map_ops.tensor_map_erase(m, k, v2.dtype) s = map_ops.tensor_map_size(m) self.assertAllEqual(s, 1) h = map_ops.tensor_map_has_key(m, k) self.assertAllEqual(h, False) l = map_ops.tensor_map_lookup(m, k2, v2.dtype) g2 = tape.gradient(l * 6, v2) self.assertAllEqual(g2, 6) del tape
def testTensorMapInsert(self): m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) v = constant_op.constant(2.0) m = map_ops.tensor_map_insert(m, k, v) s = map_ops.tensor_map_size(m) self.assertAllEqual(s, 1)
def testEraseFirstGrad(self): with backprop.GradientTape(persistent=True) as tape: m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) k2 = constant_op.constant(2.0) v = constant_op.constant(11.0) v2 = constant_op.constant(22.0) tape.watch(v) tape.watch(v2) m = map_ops.tensor_map_insert(m, k, v) l = map_ops.tensor_map_lookup(m, k, v.dtype) m = map_ops.tensor_map_insert(m, k2, v2) m, e = map_ops.tensor_map_erase(m, k, v.dtype) l2 = map_ops.tensor_map_lookup(m, k2, v2.dtype) self.assertAllClose(l2, v2) self.assertAllClose(e, v) g = tape.gradient(l * 5, v) self.assertAllEqual(g, 5) g2 = tape.gradient(l2 * 6, v2) self.assertAllEqual(g2, 6) g3 = tape.gradient(e * 7, v) self.assertAllEqual(g3, 7) m, e2 = map_ops.tensor_map_erase(m, k2, v2.dtype) g4 = tape.gradient(e2 * 8, v2) self.assertAllEqual(g4, 8) del tape
def testTensorMapLookup(self): m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) v = constant_op.constant(2.0) m = map_ops.tensor_map_insert(m, k, v) l = map_ops.tensor_map_lookup(m, k, dtypes.float32) self.assertAllClose(l, v)
def testTensorMapLookupFromEmptyMapFails(self): m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) with self.assertRaisesRegex(errors.InvalidArgumentError, "Trying to lookup non-existent key."): l = map_ops.tensor_map_lookup(m, k, dtypes.float32) self.evaluate(l)
def testStackKeysEmptyMapFails(self): m = map_ops.empty_tensor_map() with self.assertRaisesRegex( errors.InvalidArgumentError, "TensorMapStackKeys cannot be called " "on empty map."): keys = map_ops.tensor_map_stack_keys(m, dtypes.float32) self.evaluate(keys)
def testStackKeysIncorrectDtypeFails(self): m = map_ops.empty_tensor_map() k = constant_op.constant("key_with_wrong_dtype") v = constant_op.constant(2.0) m = map_ops.tensor_map_insert(m, k, v) simple = "Key does not match requested dtype." with self.assertRaisesRegex(errors.InvalidArgumentError, simple): keys = map_ops.tensor_map_stack_keys(m, dtypes.float32) self.evaluate(keys)
def testTensorMapEraseMissingKeyFails(self): m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) k2 = constant_op.constant(2.0) v = constant_op.constant(2.0) m = map_ops.tensor_map_insert(m, k2, v) with self.assertRaisesRegex(errors.InvalidArgumentError, "Trying to erase non-existent item."): m, e = map_ops.tensor_map_erase(m, k, dtypes.float32) self.evaluate(e)
def testTensorMapLookupMissingKeyFails(self): m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) k2 = constant_op.constant(2.0) v = constant_op.constant(11.0) m = map_ops.tensor_map_insert(m, k, v) with self.assertRaisesRegex(errors.InvalidArgumentError, "Trying to lookup non-existent key."): l = map_ops.tensor_map_lookup(m, k2, dtypes.float32) self.evaluate(l)
def testInsertLookupGrad(self): with backprop.GradientTape() as tape: m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) v = constant_op.constant(11.0) tape.watch(v) m = map_ops.tensor_map_insert(m, k, v) l = map_ops.tensor_map_lookup(m, k, dtypes.float32) l *= 5 g = tape.gradient(l, v) self.assertAllEqual(g, 5)
def testTensorMapHasKey(self): m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) k2 = constant_op.constant(2.0) v = constant_op.constant(2.0) m = map_ops.tensor_map_insert(m, k, v) # Check has key. b = map_ops.tensor_map_has_key(m, k) b2 = map_ops.tensor_map_has_key(m, k2) self.assertAllEqual(b, True) self.assertAllEqual(b2, False)
def testTensorMapErase(self): m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) v = constant_op.constant(2.0) m = map_ops.tensor_map_insert(m, k, v) s = map_ops.tensor_map_size(m) self.assertAllEqual(s, 1) m, e = map_ops.tensor_map_erase(m, k, dtypes.float32) s = map_ops.tensor_map_size(m) self.assertAllEqual(s, 0) self.assertAllClose(e, v)
def testStackKeysIncorrectShapeFails(self): m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) k2 = constant_op.constant([1.0, 11.0]) v = constant_op.constant(2.0) v2 = constant_op.constant(22.0) m = map_ops.tensor_map_insert(m, k, v) m = map_ops.tensor_map_insert(m, k2, v2) with self.assertRaisesRegex(errors.InvalidArgumentError, "Keys must all have the same shape."): keys = map_ops.tensor_map_stack_keys(m, dtypes.float32) self.evaluate(keys)
def testInsertLookupComposeGrad(self): with backprop.GradientTape() as tape: m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) k2 = constant_op.constant(2.0) v = constant_op.constant(11.0) tape.watch(v) m = map_ops.tensor_map_insert(m, k, v) l = map_ops.tensor_map_lookup(m, k, v.dtype) m = map_ops.tensor_map_insert(m, k2, l) l2 = map_ops.tensor_map_lookup(m, k2, l.dtype) g = tape.gradient(l2 * 5, v) self.assertAllEqual(g, 5)
def testVectorValue(self): m = map_ops.empty_tensor_map() k = constant_op.constant([1.0, 2.0]) v = constant_op.constant([11.0, 22.0]) m = map_ops.tensor_map_insert(m, k, v) s = map_ops.tensor_map_size(m) self.assertAllEqual(s, 1) l = map_ops.tensor_map_lookup(m, k, v.dtype) self.assertAllEqual(l, v) m, e = map_ops.tensor_map_erase(m, k, v.dtype) s = map_ops.tensor_map_size(m) self.assertAllEqual(s, 0) self.assertAllClose(e, v)
def testStringValue(self): m = map_ops.empty_tensor_map() k = constant_op.constant("key") v = constant_op.constant("value") k2 = constant_op.constant(1.0) v2 = constant_op.constant(2.0) m = map_ops.tensor_map_insert(m, k, v) m = map_ops.tensor_map_insert(m, k2, v2) l = map_ops.tensor_map_lookup(m, k, v.dtype) self.assertAllEqual(l, v) l2 = map_ops.tensor_map_lookup(m, k2, v2.dtype) self.assertAllClose(l2, v2) m, e = map_ops.tensor_map_erase(m, k, v.dtype) self.assertAllEqual(e, v)
def testVectorValue(self): m = map_ops.empty_tensor_map() k = constant_op.constant([1.0, 2.0]) v = constant_op.constant([11.0, 22.0]) # Test insert and lookup. m = map_ops.tensor_map_insert(m, k, v) s = map_ops.tensor_map_size(m) self.assertAllEqual(s, 1) l = map_ops.tensor_map_lookup(m, k, v.dtype) self.assertAllEqual(l, v) # Test erase and has key. m = map_ops.tensor_map_erase(m, k, v.dtype) s = map_ops.tensor_map_size(m) self.assertAllEqual(s, 0) self.assertAllEqual(map_ops.tensor_map_has_key(m, k), False)
def testSameKeyInsertLookupGrad(self): with backprop.GradientTape(persistent=True) as tape: m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) v = constant_op.constant(2.0) v2 = constant_op.constant(22.0) tape.watch(v) tape.watch(v2) m = map_ops.tensor_map_insert(m, k, v) m = map_ops.tensor_map_insert(m, k, v2) l = map_ops.tensor_map_lookup(m, k, v.dtype) g = tape.gradient(l * 5, v) g2 = tape.gradient(l * 5, v2) self.assertAllClose(g, array_ops.zeros_like(v)) self.assertAllClose(g2, 5)
def testEraseInsertComposedGrad(self): with backprop.GradientTape(persistent=True) as tape: m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) k2 = constant_op.constant(2.0) v = constant_op.constant(11.0) v2 = constant_op.constant(22.0) tape.watch(v) tape.watch(v2) m = map_ops.tensor_map_insert(m, k, v) l = map_ops.tensor_map_lookup(m, k, v.dtype) m = map_ops.tensor_map_erase(m, k, v.dtype) m = map_ops.tensor_map_insert(m, k2, l) l2 = map_ops.tensor_map_lookup(m, k2, l.dtype) g = tape.gradient(l2 * 5, v) self.assertAllEqual(g, 5) del tape
def testIfHasKeyLookup(self): m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) k2 = constant_op.constant(2.0) v = constant_op.constant(2.0) m = map_ops.tensor_map_insert(m, k, v) default_value = array_ops.zeros_like(v) l = control_flow_ops.cond( map_ops.tensor_map_has_key(m, k), lambda: map_ops.tensor_map_lookup(m, k, dtypes.float32), lambda: default_value) l2 = control_flow_ops.cond( map_ops.tensor_map_has_key(m, k2), lambda: map_ops.tensor_map_lookup(m, k, dtypes.float32), lambda: default_value) self.assertAllClose(l, v) self.assertAllClose(l2, default_value)
def testLookupMultiplyGrad(self): with backprop.GradientTape(persistent=True) as tape: k = constant_op.constant(1.0) k2 = constant_op.constant(2.0) v = constant_op.constant(11.0) v2 = constant_op.constant(22.0) tape.watch(v) tape.watch(v2) m = map_ops.empty_tensor_map() m = map_ops.tensor_map_insert(m, k, v) m = map_ops.tensor_map_insert(m, k2, v2) l1 = map_ops.tensor_map_lookup(m, k, v.dtype) l2 = map_ops.tensor_map_lookup(m, k2, v2.dtype) g = tape.gradient(l1 * l2, [v, v2]) self.assertAllClose(g, [v2, v]) g2 = tape.gradient(l1 * l1, v) self.assertAllClose(g2, 2 * v) del tape
def testStringKeyValue(self): m = map_ops.empty_tensor_map() k = constant_op.constant("key") v = constant_op.constant("value") k2 = constant_op.constant(1.0) v2 = constant_op.constant(2.0) # Test insert and lookup on string key-value pair. m = map_ops.tensor_map_insert(m, k, v) m = map_ops.tensor_map_insert(m, k2, v2) l = map_ops.tensor_map_lookup(m, k, v.dtype) self.assertAllEqual(l, v) # Test lookup on float key-value pair. l2 = map_ops.tensor_map_lookup(m, k2, v2.dtype) self.assertAllClose(l2, v2) # Test erase and has key. self.assertAllEqual(map_ops.tensor_map_has_key(m, k), True) m = map_ops.tensor_map_erase(m, k, v.dtype) self.assertAllEqual(map_ops.tensor_map_has_key(m, k), False) self.assertAllEqual(map_ops.tensor_map_has_key(m, k2), True)
def testStackKeys(self): m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) k2 = constant_op.constant(2.0) k3 = constant_op.constant(3.0) v = constant_op.constant(21.0) v2 = constant_op.constant(22.0) v3 = constant_op.constant(23.0) m = map_ops.tensor_map_insert(m, k, v) m = map_ops.tensor_map_insert(m, k2, v2) keys = map_ops.tensor_map_stack_keys(m, k.dtype) expected = constant_op.constant([1.0, 2.0]) self.assertAllClose(array_ops.shape(keys), array_ops.shape(expected)) self.assertAllClose(sort_ops.sort(keys), expected) m = map_ops.tensor_map_insert(m, k3, v3) keys = map_ops.tensor_map_stack_keys(m, k.dtype) expected = constant_op.constant([1.0, 2.0, 3.0]) self.assertAllClose(array_ops.shape(keys), array_ops.shape(expected)) self.assertAllClose(sort_ops.sort(keys), expected)
def testDiffKeySameValueGrad(self): with backprop.GradientTape(persistent=True) as tape: m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) k2 = constant_op.constant(11.0) v = constant_op.constant(2.0) v2 = constant_op.constant(22.0) tape.watch(v) tape.watch(v2) m = map_ops.tensor_map_insert(m, k, v) m = map_ops.tensor_map_insert(m, k2, v) l = map_ops.tensor_map_lookup(m, k, v.dtype) l2 = map_ops.tensor_map_lookup(m, k2, v.dtype) g = tape.gradient(l + l2, v) self.assertAllEqual(g, 2) m = map_ops.tensor_map_insert(m, k2, v2) l2 = map_ops.tensor_map_lookup(m, k2, v2.dtype) g2 = tape.gradient(l + l2, v2) self.assertAllEqual(g2, 1) del tape
def testLookupAddGrad(self): with backprop.GradientTape(persistent=True) as tape: k = constant_op.constant(1.0) k2 = constant_op.constant(2.0) v = constant_op.constant(11.0) v2 = constant_op.constant(22.0) tape.watch(v) tape.watch(v2) m = map_ops.empty_tensor_map() m = map_ops.tensor_map_insert(m, k, v) m = map_ops.tensor_map_insert(m, k2, v2) l1 = map_ops.tensor_map_lookup(m, k, v.dtype) l2 = map_ops.tensor_map_lookup(m, k2, v2.dtype) g = tape.gradient(l1 + l2, [l1, l2]) self.assertAllClose(g, [1, 1]) g2 = tape.gradient(l1 + l2, [v, v2]) self.assertAllClose(g2, [1, 1]) g3 = tape.gradient(l1 + l2 * 4, v2) self.assertAllEqual(g3, 4) del tape
def testReplaceLookupGrad(self): with backprop.GradientTape(persistent=True) as tape: m = map_ops.empty_tensor_map() k = constant_op.constant(1.0) v = constant_op.constant(11.0) v2 = constant_op.constant(22.0) tape.watch(v) tape.watch(v2) m = map_ops.tensor_map_insert(m, k, v) l = map_ops.tensor_map_lookup(m, k, v.dtype) self.assertAllClose(l, v) g = tape.gradient(l * 5, v) self.assertAllEqual(g, 5) m = map_ops.tensor_map_insert(m, k, v2) l2 = map_ops.tensor_map_lookup(m, k, v2.dtype) self.assertAllClose(l2, v2) g2 = tape.gradient(l2 * 6, v) g3 = tape.gradient(l2 * 7, v2) self.assertAllClose(g2, array_ops.zeros_like(v)) self.assertAllClose(g3, 7) del tape
def testEmptyTensorMapSize(self): m = map_ops.empty_tensor_map() s = map_ops.tensor_map_size(m) self.assertAllEqual(s, 0)