tf.data.Dataset skip() #446

Oceania2018 · Oceania2018 · commit 68df1b70f006 · 2020-08-22T18:08:23.000-05:00
diff --git a/src/TensorFlowNET.Core/Data/DatasetV2.cs b/src/TensorFlowNET.Core/Data/DatasetV2.cs
@@ -41,6 +41,9 @@ public IDatasetV2 shard(int num_shards, int index)
         public IDatasetV2 shuffle(int buffer_size, int? seed = null, bool reshuffle_each_iteration = true)
             => new ShuffleDataset(this, buffer_size, seed: seed, reshuffle_each_iteration: reshuffle_each_iteration);
 
+        public IDatasetV2 skip(int count)
+            => new SkipDataset(this, count);
+
         public IDatasetV2 optimize(string[] optimizations, string[] optimization_configs)
             => new OptimizeDataset(this, optimizations, optimization_configs: optimization_configs);
 
diff --git a/src/TensorFlowNET.Core/Data/IDatasetV2.cs b/src/TensorFlowNET.Core/Data/IDatasetV2.cs
@@ -34,6 +34,13 @@ public interface IDatasetV2 : IEnumerable<(Tensor, Tensor)>
 
         IDatasetV2 shuffle(int buffer_size, int? seed = null, bool reshuffle_each_iteration = true);
 
+        /// <summary>
+        /// Creates a `Dataset` that skips `count` elements from this dataset.
+        /// </summary>
+        /// <param name="count"></param>
+        /// <returns></returns>
+        IDatasetV2 skip(int count);
+
         IDatasetV2 batch(int batch_size, bool drop_remainder = false);
 
         IDatasetV2 prefetch(int buffer_size = -1, int? slack_period = null);
diff --git a/src/TensorFlowNET.Core/Data/SkipDataset.cs b/src/TensorFlowNET.Core/Data/SkipDataset.cs
@@ -0,0 +1,24 @@
+﻿using System;
+using System.Collections.Generic;
+using System.Text;
+using static Tensorflow.Binding;
+
+namespace Tensorflow
+{
+    /// <summary>
+    /// A `Dataset` skipping the first `count` elements from its input.
+    /// </summary>
+    public class SkipDataset : UnaryUnchangedStructureDataset
+    {
+        Tensor _count;
+
+        public SkipDataset(IDatasetV2 input_dataset,
+            int count) : base(input_dataset)
+        {
+            _count = tf.convert_to_tensor(count, dtype: dtypes.int64, name: "count");
+            variant_tensor = ops.skip_dataset(input_dataset.variant_tensor, 
+                    _count,
+                    output_types, output_shapes);
+        }
+    }
+}
diff --git a/src/TensorFlowNET.Core/Operations/dataset_ops.cs b/src/TensorFlowNET.Core/Operations/dataset_ops.cs
@@ -106,6 +106,24 @@ public Tensor shuffle_dataset_v3(Tensor input_dataset, Tensor buffer_size,
             throw new NotImplementedException("");
         }
 
+        public Tensor skip_dataset(Tensor input_dataset, Tensor count,
+            TF_DataType[] output_types, TensorShape[] output_shapes,
+            string name = null)
+        {
+            if (tf.Context.executing_eagerly())
+            {
+                var results = tf.Runner.TFE_FastPathExecute(tf.Context, tf.Context.DeviceName,
+                    "SkipDataset", name,
+                    null,
+                    input_dataset, count,
+                    "output_types", output_types,
+                    "output_shapes", output_shapes);
+                return results[0];
+            }
+
+            throw new NotImplementedException("");
+        }
+
         public Tensor dummy_seed_generator(string name = null)
         {
             if (tf.Context.executing_eagerly())
diff --git a/test/TensorFlowNET.UnitTest/Dataset/DatasetTest.cs b/test/TensorFlowNET.UnitTest/Dataset/DatasetTest.cs
@@ -84,5 +84,20 @@ public void Shard()
                 value += 3;
             }
         }
+
+        [TestMethod]
+        public void Skip()
+        {
+            long value = 7;
+
+            var dataset = tf.data.Dataset.range(10);
+            dataset = dataset.skip(7);
+
+            foreach (var item in dataset)
+            {
+                Assert.AreEqual(value, (long)item.Item1);
+                value ++;
+            }
+        }
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -84,5 +84,20 @@ public void Shard()`
`84`	`84`	`value += 3;`
`85`	`85`	`}`
`86`	`86`	`}`
	`87`	`+`
	`88`	`+ [TestMethod]`
	`89`	`+ public void Skip()`
	`90`	`+ {`
	`91`	`+ long value = 7;`
	`92`	`+`
	`93`	`+ var dataset = tf.data.Dataset.range(10);`
	`94`	`+ dataset = dataset.skip(7);`
	`95`	`+`
	`96`	`+ foreach (var item in dataset)`
	`97`	`+ {`
	`98`	`+ Assert.AreEqual(value, (long)item.Item1);`
	`99`	`+ value ++;`
	`100`	`+ }`
	`101`	`+ }`
`87`	`102`	`}`
`88`	`103`	`}`