-
Notifications
You must be signed in to change notification settings - Fork 11
Expand file tree
/
Copy pathTranspose.cs
More file actions
43 lines (42 loc) · 1.83 KB
/
Copy pathTranspose.cs
File metadata and controls
43 lines (42 loc) · 1.83 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
35
36
37
38
39
40
41
42
43
using Tensornet.Common;
using Tensornet.Native;
using Tensornet.Exceptions;
using Tensornet.Native.Param;
namespace Tensornet{
public static class TransposeExtension{
/// <summary>
/// Reverse the two axes of a tensor.
/// </summary>
/// <typeparam name="T"></typeparam>
/// <param name="src"></param>
/// <param name="dimA"></param>
/// <param name="dimB"></param>
/// <returns></returns>
public static Tensor<T> Transpose<T>(this Tensor<T> src, int dimA, int dimB) where T : struct, IEquatable<T>, IConvertible
{
Tensor<T> res = new Tensor<T>(DeduceLayout(src.TLayout, dimA, dimB));
res.TLayout.InitContiguousLayout();
TransposeInternal(src, res, dimA, dimB);
return res;
}
private unsafe static void TransposeInternal<T>(Tensor<T> src, Tensor<T> dst, int dimA, int dimB) where T : struct, IEquatable<T>, IConvertible{
TransposeParam param = new TransposeParam() { dimA = dimA, dimB = dimB };
IntPtr status = NativeExecutor.Execute(NativeApi.Transpose, src.TMemory, dst.TMemory, src.TLayout, dst.TLayout, new IntPtr(¶m), Tensor<T>.Provider);
NativeStatus.AssertOK(status);
}
private static TensorLayout DeduceLayout(TensorLayout src, int dimA, int dimB){
TensorLayout res = new TensorLayout();
if (dimA >= src.NDim || dimB >= src.NDim) {
throw new InvalidParamException("Invalid param for transpose.");
}
res.DType = src.DType;
res.NDim = src.NDim;
for (int i = 0; i < src.NDim; i++) {
res.Shape[i] = src.Shape[i];
}
res.Shape[dimA] = src.Shape[dimB];
res.Shape[dimB] = src.Shape[dimA];
return res;
}
}
}