def testCondPacked(self):
x = MaskedTensorV2([1, 2, 3, 4], [True, False, True, False])
y = MaskedTensorV2([5, 6, 7, 8], [False, True, True, False])
x = extension_type.pack(x)
y = extension_type.pack(y)
x_2 = control_flow_ops.cond(constant_op.constant(True), lambda: x,
lambda: y)
y_2 = control_flow_ops.cond(constant_op.constant(False), lambda: x,
lambda: y)
self.assertAllEqual(x.values, x_2.values)
self.assertAllEqual(x.mask, x_2.mask)
self.assertAllEqual(y.values, y_2.values)
self.assertAllEqual(y.mask, y_2.mask)
a = MaskedTensorV2([1, 2, 3, 4], [True, False, True, False])
b = extension_type.pack(a)
b = control_flow_ops.cond(constant_op.constant(True),
lambda: array_ops.size(a.mask),
lambda: array_ops.size(a.values))
self.assertAllEqual(b, 4)
# Note: the following example would fail (with `Retval[0] does not have a
# value`) if `ExtensionType.__getattr__` cached the results of unpacking
# the value. See the comment in `ExtensionType.__getattr__` for details.
c = MaskedTensorV2([1, 2, 3, 4], [True, False, True, False])
c = extension_type.pack(c)
d = control_flow_ops.cond(constant_op.constant(False),
lambda: array_ops.size(c.mask),
lambda: array_ops.size(c.values))
self.assertAllEqual(d, 4)
def testPackedEncoding(self):
mt1 = MaskedTensorV2([1, 2, 3, 4], [True, True, False, True])
self.assertLen(nest.flatten(mt1, expand_composites=True), 2)
mt2 = extension_type.pack(mt1)
self.assertLen(nest.flatten(mt2, expand_composites=True), 1)
self.assertIsInstance(mt2.values, ops.Tensor)
self.assertAllEqual(mt2.values, [1, 2, 3, 4])
self.assertIsInstance(mt2.mask, ops.Tensor)
self.assertAllEqual(mt2.mask, [True, True, False, True])
mt3 = extension_type.unpack(mt2)
self.assertLen(nest.flatten(mt3, expand_composites=True), 2)
self.assertIsInstance(mt3.values, ops.Tensor)
self.assertAllEqual(mt3.values, [1, 2, 3, 4])
self.assertIsInstance(mt3.mask, ops.Tensor)
self.assertAllEqual(mt3.mask, [True, True, False, True])
nest.assert_same_structure(mt1, mt3, expand_composites=True)
with self.assertRaisesRegex(ValueError, "don't have the same"): # pylint: disable=g-error-prone-assert-raises
nest.assert_same_structure(mt1, mt2, expand_composites=True)
mt4 = MaskedTensorV1([1, 2, 3, 4], [True, True, False, True])
with self.assertRaisesRegex(
ValueError,
'ExtensionTypes must have a __name__ field in order to be packed.'
):
extension_type.pack(mt4)
def testPassIntoTfFunction(self):
@def_function.function
def fn(x):
return x.with_default(99)
mt = MaskedTensorV2([1, 2, 3, 4], [True, True, False, True])
self.assertAllEqual([1, 2, 99, 4], fn(mt))
self.assertAllEqual([1, 2, 99, 4], fn(extension_type.pack(mt)))
def testAttributeAccessors(self):
mt1 = MaskedTensorV2([1, 2, 3, 4], [True, True, False, True])
mt2 = extension_type.pack(mt1)
for mt in [mt1, mt2]:
self.assertIsInstance(mt.values, ops.Tensor)
self.assertAllEqual(mt.values, [1, 2, 3, 4])
self.assertIsInstance(mt.mask, ops.Tensor)
self.assertAllEqual(mt.mask, [True, True, False, True])
def testWhileLoopPacked(self):
x = MaskedTensorV2([1, 2, 3, 4], [True, False, True, False])
x = extension_type.pack(x)
cond = lambda i, x: i < 10
def body(i, x):
return i + 1, extension_type.pack(
MaskedTensorV2(x.values * 2, x.mask))
_, y = control_flow_ops.while_loop_v2(cond, body, [0, x])
self.assertIsInstance(y, MaskedTensorV2)
self.assertAllEqual(y.values, [1024, 2048, 3072, 4096])
self.assertAllEqual(y.mask, [True, False, True, False])
def testAttributesAreImmutable(self):
mt1 = MaskedTensorV2([1, 2, 3, 4], [True, True, False, True])
mt2 = extension_type.pack(mt1)
for mt in [mt1, mt2]:
with self.assertRaisesRegex(AttributeError,
"cannot assign to field 'score'"):
mt.score = 12
with self.assertRaisesRegex(AttributeError,
"cannot assign to field 'values'"):
mt.values = constant_op.constant([4, 3, 2, 1])
with self.assertRaisesRegex(AttributeError,
"cannot delete field 'values'"):
del mt.values
def testAttributesAreImmutable(self):
mt1 = MaskedTensorV2([1, 2, 3, 4], [True, True, False, True])
mt2 = extension_type.pack(mt1)
for mt in [mt1, mt2]:
with self.assertRaisesRegex(
AttributeError,
'Cannot mutate attribute `score` outside the custom constructor of ExtensionType'
):
mt.score = 12
with self.assertRaisesRegex(
AttributeError,
'Cannot mutate attribute `values` outside the custom constructor of ExtensionType'
):
mt.values = constant_op.constant([4, 3, 2, 1])
with self.assertRaisesRegex(
AttributeError,
'Cannot mutate attribute `values` outside the custom constructor of ExtensionType'
):
del mt.values
def body(i, x):
return i + 1, extension_type.pack(
MaskedTensorV2(x.values * 2, x.mask))
def mask_neg_values_packed(x):
return extension_type.pack(MaskedTensorV2(x, x > 0))