def testZerosLikeObjectAlt(self, values, dtype, expected):
st = StructuredTensor.from_pyval(values)
# NOTE: zeros_like is very robust. There aren't arguments that
# should cause this operation to fail.
actual = array_ops.zeros_like(st, dtype)
self.assertAllEqual(actual, expected)
actual2 = array_ops.zeros_like_v2(st, dtype)
self.assertAllEqual(actual2, expected)
def testZerosLikeObject(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: zeros_like is very robust. There aren't arguments that
# should cause this operation to fail.
actual = array_ops.zeros_like(st, dtype)
self.assertAllEqual(actual, expected)
actual2 = array_ops.zeros_like_v2(st, dtype)
self.assertAllEqual(actual2, expected)
def step_fn(example):
segment_ids = array_ops.zeros_like_v2(example)
num_segment = array_ops.shape(example)[0]
# If number of segments is dynamic, output should be a dynamic shape.
return math_ops.unsorted_segment_sum(example, segment_ids,
num_segment)