diff --git a/src/TensorFlowNET.Core/APIs/tf.array.cs b/src/TensorFlowNET.Core/APIs/tf.array.cs
index 4d9c3da5..b529cd31 100644
--- a/src/TensorFlowNET.Core/APIs/tf.array.cs
+++ b/src/TensorFlowNET.Core/APIs/tf.array.cs
@@ -140,6 +140,16 @@ namespace Tensorflow
public Tensor gather(Tensor @params, Tensor indices, string name = null, int axis = 0)
=> array_ops.gather(@params, indices, name: name, axis: ops.convert_to_tensor(axis));
+ ///
+ /// Gather slices from `params` into a Tensor with shape specified by `indices`.
+ ///
+ ///
+ ///
+ ///
+ ///
+ public Tensor gather_nd(Tensor @params, Tensor indices, string name = null)
+ => gen_array_ops.gather_nd(@params, indices, name: name);
+
///
/// Return the elements, either from `x` or `y`, depending on the `condition`.
///
diff --git a/src/TensorFlowNET.Core/Gradients/array_grad.cs b/src/TensorFlowNET.Core/Gradients/array_grad.cs
index 016e4f02..a4da60ee 100644
--- a/src/TensorFlowNET.Core/Gradients/array_grad.cs
+++ b/src/TensorFlowNET.Core/Gradients/array_grad.cs
@@ -403,7 +403,26 @@ namespace Tensorflow.Gradients
input_grad.set_shape(op.inputs[0].GetShape());
}
return new Tensor[] { input_grad, null };
+ }
+ [RegisterGradient("GatherNd")]
+ public static Tensor[] _GatherNdGrad(Operation op, Tensor[] grads)
+ {
+ var @ref = op.inputs[0];
+ var indices = op.inputs[1];
+ var grad = grads[0];
+ var ref_shape = array_ops.shape(@ref, out_type: indices.dtype);
+ Tensor ref_grad = null;
+ if (indices.shape.ndim == 2 && indices.shape.dims[indices.shape.Length - 1] == 1)
+ {
+ ref_grad = (Tensor)new IndexedSlices(grad, array_ops.squeeze(indices, axis: -1), ref_shape);
+ }
+ else
+ {
+ ref_grad = gen_array_ops.scatter_nd(indices, grad, ref_shape);
+ }
+ return new Tensor[] { ref_grad, null };
}
+
}
}
diff --git a/src/TensorFlowNET.Core/Operations/array_ops.cs b/src/TensorFlowNET.Core/Operations/array_ops.cs
index abf44c64..57af3b83 100644
--- a/src/TensorFlowNET.Core/Operations/array_ops.cs
+++ b/src/TensorFlowNET.Core/Operations/array_ops.cs
@@ -829,7 +829,7 @@ namespace Tensorflow
/// A `Tensor`. Has the same type as `input`.
/// Contains the same data as `input`, but has one or more dimensions of
/// size 1 removed.
- public static Tensor squeeze(Tensor input, int[] axis = null, string name = null)
+ public static Tensor squeeze(Tensor input, Axis axis = null, string name = null)
=> gen_array_ops.squeeze(input, axis, name);
public static Tensor identity(Tensor input, string name = null)
diff --git a/test/TensorFlowNET.UnitTest/GradientTest/GradientEagerTest.cs b/test/TensorFlowNET.UnitTest/GradientTest/GradientEagerTest.cs
index ed759904..1cfceb3e 100644
--- a/test/TensorFlowNET.UnitTest/GradientTest/GradientEagerTest.cs
+++ b/test/TensorFlowNET.UnitTest/GradientTest/GradientEagerTest.cs
@@ -62,7 +62,7 @@ namespace TensorFlowNET.UnitTest.Gradient
// Calcute the gradient of (x1-x2)^2
// by Automatic Differentiation in Eager mode
// Expected is 2*(abs(x1-x2))
- Tensor x1 = new NDArray( new float[] { 1, 3, 5, 21, 19, 17 });
+ Tensor x1 = new NDArray(new float[] { 1, 3, 5, 21, 19, 17 });
Tensor x2 = new NDArray(new float[] { 29, 27, 23, 7, 11, 13 });
float[] expected = new float[]
{
@@ -187,5 +187,20 @@ namespace TensorFlowNET.UnitTest.Gradient
Assert.AreEqual((float)grad.numpy(), 2.0f);
}
}
+
+ [TestMethod]
+ public void GatherNdTest()
+ {
+ var x = tf.constant(new float[,] { { 1.0f, 2.0f, 3.0f }, { 1.0f, 2.0f, 3.0f }, { 1.0f, 2.0f, 3.0f } }, dtype: TF_DataType.TF_FLOAT);
+ var indices = tf.constant(new int[,] { { 0, 1 }, { 1, 1 }, { 2, 1 } }, dtype: TF_DataType.TF_INT32);
+ using (var tape = tf.GradientTape())
+ {
+ tape.watch(x);
+ var res = tf.gather_nd(x, indices);
+ var grad = tape.gradient(res, x);
+ var expected = np.array(new float[,] { { 0f, 1f, 0f }, { 0f, 1f, 0f }, { 0f, 1f, 0f } });
+ Assert.IsTrue(Enumerable.SequenceEqual(grad.ToArray(), expected.ToArray()));
+ }
+ }
}
}