Adding 'set/get_default_device()'

NiklasGustafsson · NiklasGustafsson · commit d4ff8f944514 · 2024-05-31T16:20:39.000-07:00
diff --git a/src/TorchSharp/FFT.cs b/src/TorchSharp/FFT.cs
@@ -350,6 +350,10 @@ public static Tensor ifftshift(Tensor input, long[] dim = null)
             /// <param name="requires_grad">If autograd should record operations on the returned tensor.</param>
             public static Tensor fftfreq(long n, double d = 1.0, torch.ScalarType? dtype = null, torch.Device device = null, bool requires_grad = false)
             {
+                if (device is null)
+                {
+                    device = get_default_device();
+                }
                 device = torch.InitializeDevice(device);
                 if (!dtype.HasValue) {
                     // Determine the element type dynamically.
@@ -376,6 +380,10 @@ public static Tensor fftfreq(long n, double d = 1.0, torch.ScalarType? dtype = n
             /// <param name="requires_grad">If autograd should record operations on the returned tensor.</param>
             public static Tensor rfftfreq(long n, double d = 1.0, torch.ScalarType? dtype = null, torch.Device device = null, bool requires_grad = false)
             {
+                if (device is null)
+                {
+                    device = get_default_device();
+                }
                 device = torch.InitializeDevice(device);
                 if (!dtype.HasValue) {
                     // Determine the element type dynamically.
diff --git a/src/TorchSharp/NN/Module.cs b/src/TorchSharp/NN/Module.cs
@@ -1063,8 +1063,10 @@ protected static (Device device, ScalarType dtype) GetDefaultDeviceAndType(Devic
                     if (!dtype.HasValue)
                         dtype = get_default_dtype();
 
-                    if (device == null)
-                        device = torch.CPU;
+                    if (device is null)
+                    {
+                        device = get_default_device();
+                    }
 
                     return (device, dtype.Value);
                 }
diff --git a/src/TorchSharp/Tensor/Factories/Tensor.Factories.cs b/src/TorchSharp/Tensor/Factories/Tensor.Factories.cs
@@ -28,6 +28,10 @@ public static partial class torch
         /// <param name="requires_grad"> If autograd should record operations on the returned tensor. Default: false.</param>
         public static Tensor arange(Scalar start, Scalar stop, Scalar step, ScalarType? dtype = null, Device? device = null, bool requires_grad = false)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
 
             if (!dtype.HasValue) {
@@ -77,6 +81,10 @@ public static Tensor arange(Scalar stop, ScalarType? dtype = null, Device? devic
         /// </summary>
         public static Tensor eye(long rows, long columns = -1L, ScalarType? dtype = null, Device? device = null, bool requires_grad = false, string[]? names = null)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
@@ -440,6 +448,10 @@ public static Tensor polar(Tensor abs, Tensor angle)
 
         public static Tensor from_file(string filename, bool? shared = null, long? size = 0, ScalarType? dtype = null, Device? device = null, bool requires_grad = false)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
@@ -456,6 +468,10 @@ public static Tensor from_file(string filename, bool? shared = null, long? size
         /// </summary>
         public static Tensor linspace(double start, double end, long steps, ScalarType? dtype = null, Device? device = null, bool requires_grad = false)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
@@ -477,6 +493,10 @@ public static Tensor linspace(double start, double end, long steps, ScalarType?
         /// </summary>
         public static Tensor logspace(double start, double end, long steps, double @base = 10, ScalarType? dtype = null, Device? device = null, bool requires_grad = false)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
@@ -579,6 +599,7 @@ private static void ValidateIntegerRange(long value, ScalarType dtype, string ar
 
         private static ConcurrentDictionary<TorchSharp.PInvoke.GCHandleDeleter, TorchSharp.PInvoke.GCHandleDeleter> deleters;
         private static ScalarType default_dtype = ScalarType.Float32;
+        private static Device default_device = CPU;
 
         static torch()
         {
diff --git a/src/TorchSharp/Tensor/Factories/empty.cs b/src/TorchSharp/Tensor/Factories/empty.cs
@@ -98,6 +98,10 @@ public static Tensor empty(int dim0, int dim1, int dim2, int dim3, ScalarType? d
         /// </summary>
         public static Tensor empty_strided(long[] size, long[] strides, ScalarType? dtype = null, Device? device = null, bool requires_grad = false, string[]? names = null)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
@@ -130,6 +134,10 @@ public static Tensor empty_strided(long[] size, long[] strides, ScalarType? dtyp
         /// </summary>
         private static Tensor _empty(ReadOnlySpan<long> size, ScalarType? dtype = null, Device? device = null, bool requires_grad = false, string[]? names = null)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
diff --git a/src/TorchSharp/Tensor/Factories/full.cs b/src/TorchSharp/Tensor/Factories/full.cs
@@ -97,6 +97,10 @@ public static Tensor full(int dim0, int dim1, int dim2, int dim3, Scalar value,
         /// </summary>
         private static Tensor _full(ReadOnlySpan<long> size, Scalar value, ScalarType? dtype = null, Device? device = null, bool requires_grad = false, string[]? names = null)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
diff --git a/src/TorchSharp/Tensor/Factories/ones.cs b/src/TorchSharp/Tensor/Factories/ones.cs
@@ -97,6 +97,10 @@ public static Tensor ones(int dim0, int dim1, int dim2, int dim3, ScalarType? dt
         /// </summary>
         private static Tensor _ones(ReadOnlySpan<long> size, ScalarType? dtype = null, Device? device = null, bool requires_grad = false, string[]? names = null)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
diff --git a/src/TorchSharp/Tensor/Factories/rand.cs b/src/TorchSharp/Tensor/Factories/rand.cs
@@ -348,6 +348,10 @@ private static Tensor _rand(ReadOnlySpan<long> size, ScalarType? dtype = null, D
             if (dtype.HasValue && torch.is_integral(dtype.Value))
                 throw new ArgumentException($"torch.rand() was passed a bad dtype: {dtype}. It must be floating point or complex.", "dtype");
 
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
@@ -476,6 +480,10 @@ private static Tensor _randn(ReadOnlySpan<long> size, ScalarType? dtype = null,
             if (dtype.HasValue && torch.is_integral(dtype.Value))
                 throw new ArgumentException($"torch.randn() was passed a bad dtype: {dtype}. It must be floating point or complex.", "dtype");
 
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
diff --git a/src/TorchSharp/Tensor/Factories/zeros.cs b/src/TorchSharp/Tensor/Factories/zeros.cs
@@ -94,6 +94,10 @@ public static Tensor zeros(int dim0, int dim1, int dim2, int dim3, ScalarType? d
 
         private static Tensor _zeros(ReadOnlySpan<long> size, ScalarType? dtype = null, Device? device = null, bool requires_grad = false, string[]? names = null)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
diff --git a/src/TorchSharp/Tensor/torch.SpectralOps.cs b/src/TorchSharp/Tensor/torch.SpectralOps.cs
@@ -51,6 +51,10 @@ public static Tensor istft(Tensor input, long n_fft, long hop_length = -1, long
         /// </summary>
         public static Tensor bartlett_window(long len, bool periodic = true, ScalarType? dtype = null, Device? device = null, bool requires_grad = false)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
@@ -74,6 +78,10 @@ public static Tensor bartlett_window(long len, bool periodic = true, ScalarType?
         /// </summary>
         public static Tensor blackman_window(long len, bool periodic = true, ScalarType? dtype = null, Device? device = null, bool requires_grad = false)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
@@ -98,6 +106,10 @@ public static Tensor blackman_window(long len, bool periodic = true, ScalarType?
         /// </summary>
         public static Tensor hamming_window(long len, bool periodic = true, float alpha = 0.54f, float beta = 0.46f, ScalarType? dtype = null, Device? device = null, bool requires_grad = false)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
@@ -121,6 +133,10 @@ public static Tensor hamming_window(long len, bool periodic = true, float alpha
         /// </summary>
         public static Tensor hann_window(long len, bool periodic = true, ScalarType? dtype = null, Device? device = null, bool requires_grad = false)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
@@ -144,6 +160,10 @@ public static Tensor hann_window(long len, bool periodic = true, ScalarType? dty
         /// </summary>
         public static Tensor kaiser_window(long len, bool periodic = true, float beta = 12.0f, ScalarType? dtype = null, Device? device = null, bool requires_grad = false)
         {
+            if (device is null)
+            {
+                device = get_default_device();
+            }
             device = InitializeDevice(device);
             if (!dtype.HasValue) {
                 // Determine the element type dynamically.
diff --git a/src/TorchSharp/Tensor/torch.Tensors.cs b/src/TorchSharp/Tensor/torch.Tensors.cs
@@ -61,9 +61,23 @@ public static void set_default_dtype(ScalarType dtype)
         /// <returns></returns>
         [Pure]public static ScalarType get_default_dtype() => default_dtype;
 
+        [Pure]public static Device get_default_device() => default_device;
+
         // https://pytorch.org/docs/stable/generated/torch.set_default_tensor_type
         public static void set_default_tensor_type(Tensor t) => set_default_dtype(t.dtype);
 
+        public static void set_default_device(Device device) 
+        {
+            if (device == null) 
+                throw new ArgumentNullException(nameof(device));
+            default_device = device;
+        }
+
+        public static void set_default_device(string device)
+        {
+            set_default_device(new Device(device));
+        }
+
         // https://pytorch.org/docs/stable/generated/torch.numel
         /// <summary>
         /// Get the number of elements in the input tensor.
diff --git a/test/TorchSharpTest/TestTorchTensor.cs b/test/TorchSharpTest/TestTorchTensor.cs
@@ -8375,5 +8375,67 @@ public void TestSaveAndLoadModuleWithLarger2GBTensorCUDA()
                 }
             }
         }
+
+        [Fact]
+        public void DefaultDTypeCreation()
+        {
+            var dt = torch.get_default_dtype();
+
+            var t = torch.zeros(5,5);
+            Assert.Equal(torch.float32, t.dtype);
+
+            try {
+                torch.set_default_dtype(torch.float64);              
+                
+                t = torch.zeros(5,5);
+                Assert.Equal(torch.float64, t.dtype);
+
+                t = torch.ones(5,5);
+                Assert.Equal(torch.float64, t.dtype);
+
+                t = torch.rand(5,5);
+                Assert.Equal(torch.float64, t.dtype);
+
+                t = torch.randn(5,5);
+                Assert.Equal(torch.float64, t.dtype);
+
+                t = torch.logspace(5, 15, 20);
+                Assert.Equal(torch.float64, t.dtype);
+            }
+            finally {
+                torch.set_default_dtype(dt);
+            }
+        }
+
+        [Fact]
+        public void DefaultDeviceCreation()
+        {
+            var dt = torch.get_default_device();
+
+            var t = torch.zeros(5,5);
+            Assert.Equal(DeviceType.CPU, t.device_type);
+
+            try {
+                torch.set_default_device(torch.META);              
+                
+                t = torch.zeros(5,5);
+                Assert.Equal(DeviceType.META, t.device_type);
+
+                t = torch.ones(5,5);
+                Assert.Equal(DeviceType.META, t.device_type);
+
+                t = torch.rand(5,5);
+                Assert.Equal(DeviceType.META, t.device_type);
+
+                t = torch.randn(5,5);
+                Assert.Equal(DeviceType.META, t.device_type);
+
+                t = torch.logspace(5, 15, 20);
+                Assert.Equal(DeviceType.META, t.device_type);
+            }
+            finally {
+                torch.set_default_device(dt);
+            }
+        }
     }
 }
