From c19e89ddfcb0d96fb3de44e193c658a526920b98 Mon Sep 17 00:00:00 2001
From: Jacob G-W <jacoblevgw@gmail.com>
Date: Sun, 11 Feb 2024 17:11:35 -0500
Subject: [PATCH 1/2] add mixed precision support to deepxde

---
 deepxde/config.py | 16 +++++++++++++++-
 deepxde/model.py  |  7 +++++--
 deepxde/real.py   |  4 ++++
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/deepxde/config.py b/deepxde/config.py
index 6d87e67b6..d593e964e 100644
--- a/deepxde/config.py
+++ b/deepxde/config.py
@@ -74,7 +74,7 @@ def set_default_float(value):
     The default floating point type is 'float32'.
 
     Args:
-        value (String): 'float16', 'float32', or 'float64'.
+        value (String): 'float16', 'float32', 'float64', or 'mixed' (mixed precision in https://arxiv.org/abs/2401.16645).
     """
     if value == "float16":
         print("Set the default float type to float16")
@@ -85,6 +85,20 @@ def set_default_float(value):
     elif value == "float64":
         print("Set the default float type to float64")
         real.set_float64()
+    elif value == "mixed":
+        print("Set the float type to mixed precision of float16 and float32")
+        real.set_mixed()
+        if backend_name == "tensorflow":
+            real.set_float16()
+            tf.keras.mixed_precision.set_global_policy("mixed_float16")
+            return # don't try to set it again below
+        if backend_name == "pytorch":
+            # Use float16 during the forward and backward passes, but store in float32
+            real.set_float32()
+        else:
+            raise ValueError(
+                f"{backend_name} backend does not currently support mixed precision."
+            )
     else:
         raise ValueError(f"{value} not supported in deepXDE")
     if backend_name in ["tensorflow.compat.v1", "tensorflow"]:
diff --git a/deepxde/model.py b/deepxde/model.py
index 48ac2ef46..0078d8135 100644
--- a/deepxde/model.py
+++ b/deepxde/model.py
@@ -357,9 +357,12 @@ def closure():
                 total_loss = torch.sum(losses)
                 self.opt.zero_grad()
                 total_loss.backward()
-                return total_loss
 
-            self.opt.step(closure)
+            def closure_mixed():
+                with torch.autocast(device_type="cuda", dtype=torch.float16):
+                    closure()
+
+            self.opt.step(closure if not config.real.mixed else closure_mixed)
             if self.lr_scheduler is not None:
                 self.lr_scheduler.step()
 
diff --git a/deepxde/real.py b/deepxde/real.py
index 1ceb2fd7a..d7268de6e 100644
--- a/deepxde/real.py
+++ b/deepxde/real.py
@@ -7,6 +7,7 @@ class Real:
     def __init__(self, precision):
         self.precision = None
         self.reals = None
+        self.mixed = False
         if precision == 16:
             self.set_float16()
         elif precision == 32:
@@ -28,3 +29,6 @@ def set_float32(self):
     def set_float64(self):
         self.precision = 64
         self.reals = {np: np.float64, bkd.lib: bkd.float64}
+
+    def set_mixed(self):
+        self.mixed = True

From d97113eeecccd536883866c71ff617e8a02031ed Mon Sep 17 00:00:00 2001
From: Jacob G-W <jacoblevgw@gmail.com>
Date: Sun, 11 Feb 2024 17:16:52 -0500
Subject: [PATCH 2/2] update faq to include mixed precision

---
 docs/user/faq.rst | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/docs/user/faq.rst b/docs/user/faq.rst
index aaaa1ac0e..e0b88898f 100644
--- a/docs/user/faq.rst
+++ b/docs/user/faq.rst
@@ -10,6 +10,8 @@ General usage
   | **A**: `#5`_
 - | **Q**: By default, DeepXDE uses ``float32``. How can I use ``float64``?
   | **A**: `#28`_
+- | **Q**: How can I use mixed precision training?
+  | **A**: Use ``dde.config.set_default_float("mixed")`` with the ``tensorflow`` or ``pytorch`` backends. See `this paper <https://doi.org/10.1016/j.cma.2024.117093>`_ for more information.
 - | **Q**: I want to set the global random seeds.
   | **A**: `#353`_
 - | **Q**: GPU.