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)