forked from SciSharp/Tensor.NET
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathTensor.BoolIndex.cs
More file actions
34 lines (33 loc) · 1.54 KB
/
Copy pathTensor.BoolIndex.cs
File metadata and controls
34 lines (33 loc) · 1.54 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
using Tensornet.Native;
using Tensornet.Exceptions;
using Tensornet.Common;
namespace Tensornet{
public partial class Tensor<T> where T : struct, IEquatable<T>, IConvertible{
public Tensor<T> this[Tensor<bool> index]{
get{
Tensor<T> res = new Tensor<T>(new TensorLayout(TLayout, true));
BoolIndexInternal(this, res, index);
return res;
}
set{
if(value.TLayout.CanBroadCastTo(this.TLayout)){
var refer = new Tensor<T>(value.TMemory, value.TLayout.Broadcast(this.TLayout));
InplaceOperation.Execute<T>(this, index, refer);
}
else{
throw new MismatchedShapeException($"Cannot broadcast from {value.TLayout as TensorShape} to {TLayout as TensorShape}.");
}
}
}
private static unsafe void BoolIndexInternal(Tensor<T> src, Tensor<T> dst, Tensor<bool> index){
if(index.TLayout.CanBroadCastTo(dst.TLayout)){
var indexLayout = index.TLayout.Broadcast(dst.TLayout);
IntPtr status = NativeExecutor.Execute(NativeApi.BoolIndex, src.TMemory, index.TMemory, dst.TMemory, src.TLayout, indexLayout, dst.TLayout, IntPtr.Zero, Tensor<T>.Provider);
NativeStatus.AssertOK(status);
}
else{
throw new MismatchedShapeException($"Cannot broadcast from {(index.TLayout as TensorShape)} to {(dst.TLayout as TensorShape)}.");
}
}
}
}