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)