def testSizeAlt(self, values, dtype, expected):
st = StructuredTensor.from_pyval(values)
# NOTE: size is very robust. There aren't arguments that
# should cause this operation to fail.
actual = array_ops.size(st, out_type=dtype)
self.assertAllEqual(actual, expected)
actual2 = array_ops.size_v2(st, out_type=dtype)
self.assertAllEqual(actual2, expected)
def testSizeObject(self, row_partitions, shape, dtype, expected):
if row_partitions is not None:
row_partitions = [
row_partition.RowPartition.from_row_splits(r) for r in row_partitions
]
st = StructuredTensor.from_fields({},
shape=shape,
row_partitions=row_partitions)
# NOTE: size is very robust. There aren't arguments that
# should cause this operation to fail.
actual = array_ops.size(st, out_type=dtype)
self.assertAllEqual(actual, expected)
actual2 = array_ops.size_v2(st, out_type=dtype)
self.assertAllEqual(actual2, expected)
def computation(x):
return array_ops.size_v2(x)
