diff --git a/.devcontainer/cpu/devcontainer.json b/.devcontainer/cpu/devcontainer.json
index 55c62524..c71bc569 100644
--- a/.devcontainer/cpu/devcontainer.json
+++ b/.devcontainer/cpu/devcontainer.json
@@ -6,8 +6,8 @@
     "context": "../..",
     "dockerfile": "../Dockerfile",
     "args": {
-      "CLANG_VERSION": ""
-    }
+      "CLANG_VERSION": "",
+    },
   },
 
   // Use 'forwardPorts' to make a list of ports inside the container available locally.
@@ -26,10 +26,10 @@
         "ms-python.python",
         "ms-vsliveshare.vsliveshare",
         "DavidAnson.vscode-markdownlint",
-        "GitHub.copilot"
-      ]
-    }
-  }
+        "GitHub.copilot",
+      ],
+    },
+  },
 
   // Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root.
   // "remoteUser": "root"
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 2d9514af..74816aad 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -56,14 +56,14 @@ repos:
         args: ["--fix", "--show-fixes"]
 
   # 2023-12-11: Not use mypy for now.
-  # - repo: https://github.com/pre-commit/mirrors-mypy
-  #   rev: "v1.8.0"
-  #   hooks:
-  #     - id: mypy
-  #       files: src|tests
-  #       args: []
-  #       additional_dependencies:
-  #         - pytest
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: "v1.8.0"
+    hooks:
+      - id: mypy
+        files: src
+        args: ["--ignore-missing-imports"]
+        additional_dependencies:
+          - pytest
 
   - repo: https://github.com/codespell-project/codespell
     rev: "v2.2.6"
diff --git a/src/caustics/cosmology.py b/src/caustics/cosmology.py
index 51cee50f..29867136 100644
--- a/src/caustics/cosmology.py
+++ b/src/caustics/cosmology.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="operator"
 from abc import abstractmethod
 from math import pi
 from typing import Optional
@@ -56,7 +57,7 @@ class Cosmology(Parametrized):
         Name of the cosmological model.
     """
 
-    def __init__(self, name: str = None):
+    def __init__(self, name: Optional[str] = None):
         """
         Initialize the Cosmology.
 
@@ -296,7 +297,7 @@ def __init__(
         h0: Optional[Tensor] = torch.tensor(h0_default),
         critical_density_0: Optional[Tensor] = torch.tensor(critical_density_0_default),
         Om0: Optional[Tensor] = torch.tensor(Om0_default),
-        name: str = None,
+        name: Optional[str] = None,
     ):
         """
         Initialize a new instance of the FlatLambdaCDM class.
@@ -358,9 +359,9 @@ def critical_density(
         z: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        h0: Tensor = None,
-        critical_density_0: Tensor = None,
-        Om0: Tensor = None,
+        h0: Optional[Tensor] = None,
+        critical_density_0: Optional[Tensor] = None,
+        Om0: Optional[Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
         """
@@ -410,9 +411,9 @@ def comoving_distance(
         z: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        h0: Tensor = None,
-        critical_density_0: Tensor = None,
-        Om0: Tensor = None,
+        h0: Optional[Tensor] = None,
+        critical_density_0: Optional[Tensor] = None,
+        Om0: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -443,9 +444,9 @@ def transverse_comoving_distance(
         z: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        h0: Tensor = None,
-        critical_density_0: Tensor = None,
-        Om0: Tensor = None,
+        h0: Optional[Tensor] = None,
+        critical_density_0: Optional[Tensor] = None,
+        Om0: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         return self.comoving_distance(z, params, **kwargs)
diff --git a/src/caustics/lenses/base.py b/src/caustics/lenses/base.py
index 9232e17e..738c4407 100644
--- a/src/caustics/lenses/base.py
+++ b/src/caustics/lenses/base.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="call-overload"
 from abc import abstractmethod
 from typing import Optional, Union
 from functools import partial
@@ -21,7 +22,7 @@ class Lens(Parametrized):
     Base class for all lenses
     """
 
-    def __init__(self, cosmology: Cosmology, name: str = None):
+    def __init__(self, cosmology: Cosmology, name: Optional[str] = None):
         """
         Initializes a new instance of the Lens class.
 
@@ -609,7 +610,7 @@ def __init__(
         self,
         cosmology: Cosmology,
         z_l: Optional[Union[Tensor, float]] = None,
-        name: str = None,
+        name: Optional[str] = None,
     ):
         super().__init__(cosmology=cosmology, name=name)
         self.add_param("z_l", z_l)
@@ -622,7 +623,7 @@ def reduced_deflection_angle(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
+        z_l: Optional[Tensor] = None,
         **kwargs,
     ) -> tuple[Tensor, Tensor]:
         """
@@ -662,7 +663,7 @@ def physical_deflection_angle(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
+        z_l: Optional[Tensor] = None,
         **kwargs,
     ) -> tuple[Tensor, Tensor]:
         """
@@ -766,7 +767,7 @@ def surface_density(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
+        z_l: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -844,7 +845,7 @@ def time_delay(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
+        z_l: Optional[Tensor] = None,
         shapiro_time_delay: bool = True,
         geometric_time_delay: bool = True,
         **kwargs,
diff --git a/src/caustics/lenses/epl.py b/src/caustics/lenses/epl.py
index ab177b54..d4b09577 100644
--- a/src/caustics/lenses/epl.py
+++ b/src/caustics/lenses/epl.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="operator"
 from typing import Optional, Union
 
 import torch
@@ -80,7 +81,7 @@ def __init__(
         t: Optional[Union[Tensor, float]] = None,
         s: float = 0.0,
         n_iter: int = 18,
-        name: str = None,
+        name: Optional[str] = None,
     ):
         """
         Initialize an EPL lens model.
@@ -137,13 +138,13 @@ def reduced_deflection_angle(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        q: Tensor = None,
-        phi: Tensor = None,
-        b: Tensor = None,
-        t: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        q: Optional[Tensor] = None,
+        phi: Optional[Tensor] = None,
+        b: Optional[Tensor] = None,
+        t: Optional[Tensor] = None,
         **kwargs,
     ) -> tuple[Tensor, Tensor]:
         """
@@ -220,13 +221,13 @@ def potential(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        q: Tensor = None,
-        phi: Tensor = None,
-        b: Tensor = None,
-        t: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        q: Optional[Tensor] = None,
+        phi: Optional[Tensor] = None,
+        b: Optional[Tensor] = None,
+        t: Optional[Tensor] = None,
         **kwargs,
     ):
         """
@@ -251,7 +252,7 @@ def potential(
         ax, ay = self.reduced_deflection_angle(x, y, z_s, params)
         ax, ay = derotate(ax, ay, -phi)
         x, y = translate_rotate(x, y, x0, y0, phi)
-        return (x * ax + y * ay) / (2 - t)
+        return (x * ax + y * ay) / (2 - t)  # fmt: skip
 
     @unpack
     def convergence(
@@ -261,15 +262,15 @@ def convergence(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        q: Tensor = None,
-        phi: Tensor = None,
-        b: Tensor = None,
-        t: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        q: Optional[Tensor] = None,
+        phi: Optional[Tensor] = None,
+        b: Optional[Tensor] = None,
+        t: Optional[Tensor] = None,
         **kwargs,
-    ):
+    ) -> Tensor:
         """
         Compute the convergence of the lens, which describes the local density of the lens.
 
diff --git a/src/caustics/lenses/external_shear.py b/src/caustics/lenses/external_shear.py
index 6aee98a7..3f4c226a 100644
--- a/src/caustics/lenses/external_shear.py
+++ b/src/caustics/lenses/external_shear.py
@@ -50,7 +50,7 @@ def __init__(
         gamma_1: Optional[Union[Tensor, float]] = None,
         gamma_2: Optional[Union[Tensor, float]] = None,
         s: float = 0.0,
-        name: str = None,
+        name: Optional[str] = None,
     ):
         super().__init__(cosmology, z_l, name=name)
 
@@ -68,11 +68,11 @@ def reduced_deflection_angle(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        gamma_1: Tensor = None,
-        gamma_2: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        gamma_1: Optional[Tensor] = None,
+        gamma_2: Optional[Tensor] = None,
         **kwargs,
     ) -> tuple[Tensor, Tensor]:
         """
@@ -113,11 +113,11 @@ def potential(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        gamma_1: Tensor = None,
-        gamma_2: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        gamma_1: Optional[Tensor] = None,
+        gamma_2: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -151,11 +151,11 @@ def convergence(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        gamma_1: Tensor = None,
-        gamma_2: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        gamma_1: Optional[Tensor] = None,
+        gamma_2: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
diff --git a/src/caustics/lenses/mass_sheet.py b/src/caustics/lenses/mass_sheet.py
index d51c5a0c..612e5e2d 100644
--- a/src/caustics/lenses/mass_sheet.py
+++ b/src/caustics/lenses/mass_sheet.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="operator"
 from typing import Optional, Union
 
 import torch
@@ -48,7 +49,7 @@ def __init__(
         x0: Optional[Union[Tensor, float]] = None,
         y0: Optional[Union[Tensor, float]] = None,
         surface_density: Optional[Union[Tensor, float]] = None,
-        name: str = None,
+        name: Optional[str] = None,
     ):
         super().__init__(cosmology, z_l, name=name)
 
@@ -64,10 +65,10 @@ def reduced_deflection_angle(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        surface_density: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        surface_density: Optional[Tensor] = None,
         **kwargs,
     ) -> tuple[Tensor, Tensor]:
         """
@@ -103,10 +104,10 @@ def potential(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        surface_density: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        surface_density: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         # Meneghetti eq 3.81
@@ -120,10 +121,10 @@ def convergence(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        surface_density: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        surface_density: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         # Essentially by definition
diff --git a/src/caustics/lenses/multiplane.py b/src/caustics/lenses/multiplane.py
index 8ae1b641..eff1c12a 100644
--- a/src/caustics/lenses/multiplane.py
+++ b/src/caustics/lenses/multiplane.py
@@ -32,7 +32,9 @@ class Multiplane(ThickLens):
         List of thin lenses.
     """
 
-    def __init__(self, cosmology: Cosmology, lenses: list[ThinLens], name: str = None):
+    def __init__(
+        self, cosmology: Cosmology, lenses: list[ThinLens], name: Optional[str] = None
+    ):
         super().__init__(cosmology, name=name)
         self.lenses = lenses
         for lens in lenses:
diff --git a/src/caustics/lenses/nfw.py b/src/caustics/lenses/nfw.py
index cfe68f18..148ac1ae 100644
--- a/src/caustics/lenses/nfw.py
+++ b/src/caustics/lenses/nfw.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="operator,union-attr"
 from math import pi
 from typing import Optional, Union
 
@@ -82,7 +83,7 @@ def __init__(
         c: Optional[Union[Tensor, float]] = None,
         s: float = 0.0,
         use_case="batchable",
-        name: str = None,
+        name: Optional[str] = None,
     ):
         """
         Initialize an instance of the NFW lens class.
@@ -133,11 +134,11 @@ def get_scale_radius(
         self,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        m: Tensor = None,
-        c: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        m: Optional[Tensor] = None,
+        c: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -168,11 +169,11 @@ def get_scale_density(
         self,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        m: Tensor = None,
-        c: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        m: Optional[Tensor] = None,
+        c: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -201,11 +202,11 @@ def get_convergence_s(
         z_s,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        m: Tensor = None,
-        c: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        m: Optional[Tensor] = None,
+        c: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -388,11 +389,11 @@ def reduced_deflection_angle(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        m: Tensor = None,
-        c: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        m: Optional[Tensor] = None,
+        c: Optional[Tensor] = None,
         **kwargs,
     ) -> tuple[Tensor, Tensor]:
         """
@@ -437,11 +438,11 @@ def convergence(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        m: Tensor = None,
-        c: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        m: Optional[Tensor] = None,
+        c: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -480,11 +481,11 @@ def potential(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        m: Tensor = None,
-        c: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        m: Optional[Tensor] = None,
+        c: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
diff --git a/src/caustics/lenses/pixelated_convergence.py b/src/caustics/lenses/pixelated_convergence.py
index 10cda750..af2b84f5 100644
--- a/src/caustics/lenses/pixelated_convergence.py
+++ b/src/caustics/lenses/pixelated_convergence.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="index"
 from math import pi
 from typing import Optional
 
@@ -36,7 +37,7 @@ def __init__(
         convolution_mode: str = "fft",
         use_next_fast_len: bool = True,
         padding: str = "zero",
-        name: str = None,
+        name: Optional[str] = None,
     ):
         """Strong lensing with user provided kappa map
 
@@ -260,10 +261,10 @@ def reduced_deflection_angle(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        convergence_map: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        convergence_map: Optional[Tensor] = None,
         **kwargs,
     ) -> tuple[Tensor, Tensor]:
         """
@@ -367,10 +368,10 @@ def potential(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        convergence_map: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        convergence_map: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -453,10 +454,10 @@ def convergence(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        convergence_map: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        convergence_map: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
diff --git a/src/caustics/lenses/point.py b/src/caustics/lenses/point.py
index b640f84d..619a1482 100644
--- a/src/caustics/lenses/point.py
+++ b/src/caustics/lenses/point.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="operator"
 from typing import Optional, Union
 
 import torch
@@ -48,7 +49,7 @@ def __init__(
         y0: Optional[Union[Tensor, float]] = None,
         th_ein: Optional[Union[Tensor, float]] = None,
         s: float = 0.0,
-        name: str = None,
+        name: Optional[str] = None,
     ):
         """
         Initialize the Point class.
@@ -85,10 +86,10 @@ def reduced_deflection_angle(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        th_ein: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        th_ein: Optional[Tensor] = None,
         **kwargs,
     ) -> tuple[Tensor, Tensor]:
         """
@@ -124,10 +125,10 @@ def potential(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        th_ein: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        th_ein: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -161,10 +162,10 @@ def convergence(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        th_ein: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        th_ein: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
diff --git a/src/caustics/lenses/pseudo_jaffe.py b/src/caustics/lenses/pseudo_jaffe.py
index 671333bf..ebf1cb3b 100644
--- a/src/caustics/lenses/pseudo_jaffe.py
+++ b/src/caustics/lenses/pseudo_jaffe.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="operator"
 from math import pi
 from typing import Optional, Union
 
@@ -60,7 +61,7 @@ def __init__(
         core_radius: Optional[Union[Tensor, float]] = None,
         scale_radius: Optional[Union[Tensor, float]] = None,
         s: float = 0.0,
-        name: str = None,
+        name: Optional[str] = None,
     ):
         """
         Initialize the PseudoJaffe class.
@@ -101,12 +102,12 @@ def get_convergence_0(
         z_s,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        mass: Tensor = None,
-        core_radius: Tensor = None,
-        scale_radius: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        mass: Optional[Tensor] = None,
+        core_radius: Optional[Tensor] = None,
+        scale_radius: Optional[Tensor] = None,
         **kwargs,
     ):
         d_l = self.cosmology.angular_diameter_distance(z_l, params)
@@ -120,12 +121,12 @@ def mass_enclosed_2d(
         z_s,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        mass: Tensor = None,
-        core_radius: Tensor = None,
-        scale_radius: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        mass: Optional[Tensor] = None,
+        core_radius: Optional[Tensor] = None,
+        scale_radius: Optional[Tensor] = None,
         **kwargs,
     ):
         """
@@ -199,12 +200,12 @@ def reduced_deflection_angle(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        mass: Tensor = None,
-        core_radius: Tensor = None,
-        scale_radius: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        mass: Optional[Tensor] = None,
+        core_radius: Optional[Tensor] = None,
+        scale_radius: Optional[Tensor] = None,
         **kwargs,
     ) -> tuple[Tensor, Tensor]:
         """Calculate the deflection angle.
@@ -241,12 +242,12 @@ def potential(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        mass: Tensor = None,
-        core_radius: Tensor = None,
-        scale_radius: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        mass: Optional[Tensor] = None,
+        core_radius: Optional[Tensor] = None,
+        scale_radius: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -298,12 +299,12 @@ def convergence(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        mass: Tensor = None,
-        core_radius: Tensor = None,
-        scale_radius: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        mass: Optional[Tensor] = None,
+        core_radius: Optional[Tensor] = None,
+        scale_radius: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
diff --git a/src/caustics/lenses/sie.py b/src/caustics/lenses/sie.py
index 7ea2a5e7..85f6dde9 100644
--- a/src/caustics/lenses/sie.py
+++ b/src/caustics/lenses/sie.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="operator,union-attr"
 from typing import Optional, Union
 
 from torch import Tensor
@@ -56,7 +57,7 @@ def __init__(
         phi: Optional[Union[Tensor, float]] = None,
         b: Optional[Union[Tensor, float]] = None,
         s: float = 0.0,
-        name: str = None,
+        name: Optional[str] = None,
     ):
         """
         Initialize the SIE lens model.
@@ -98,12 +99,12 @@ def reduced_deflection_angle(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        q: Tensor = None,
-        phi: Tensor = None,
-        b: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        q: Optional[Tensor] = None,
+        phi: Optional[Tensor] = None,
+        b: Optional[Tensor] = None,
         **kwargs,
     ) -> tuple[Tensor, Tensor]:
         """
@@ -141,12 +142,12 @@ def potential(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        x0: Tensor = None,
-        z_l: Tensor = None,
-        y0: Tensor = None,
-        q: Tensor = None,
-        phi: Tensor = None,
-        b: Tensor = None,
+        x0: Optional[Tensor] = None,
+        z_l: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        q: Optional[Tensor] = None,
+        phi: Optional[Tensor] = None,
+        b: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -181,12 +182,12 @@ def convergence(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        q: Tensor = None,
-        phi: Tensor = None,
-        b: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        q: Optional[Tensor] = None,
+        phi: Optional[Tensor] = None,
+        b: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
diff --git a/src/caustics/lenses/singleplane.py b/src/caustics/lenses/singleplane.py
index 9644f7d2..f9b858c1 100644
--- a/src/caustics/lenses/singleplane.py
+++ b/src/caustics/lenses/singleplane.py
@@ -27,7 +27,11 @@ class SinglePlane(ThinLens):
     """
 
     def __init__(
-        self, cosmology: Cosmology, lenses: list[ThinLens], name: str = None, **kwargs
+        self,
+        cosmology: Cosmology,
+        lenses: list[ThinLens],
+        name: Optional[str] = None,
+        **kwargs,
     ):
         """
         Initialize the SinglePlane lens model.
diff --git a/src/caustics/lenses/sis.py b/src/caustics/lenses/sis.py
index 15a54172..df4f4d76 100644
--- a/src/caustics/lenses/sis.py
+++ b/src/caustics/lenses/sis.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="operator"
 from typing import Optional, Union
 
 from torch import Tensor
@@ -47,7 +48,7 @@ def __init__(
         y0: Optional[Union[Tensor, float]] = None,
         th_ein: Optional[Union[Tensor, float]] = None,
         s: float = 0.0,
-        name: str = None,
+        name: Optional[str] = None,
     ):
         """
         Initialize the SIS lens model.
@@ -67,10 +68,10 @@ def reduced_deflection_angle(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        th_ein: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        th_ein: Optional[Tensor] = None,
         **kwargs,
     ) -> tuple[Tensor, Tensor]:
         """
@@ -106,10 +107,10 @@ def potential(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        th_ein: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        th_ein: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -143,10 +144,10 @@ def convergence(
         z_s: Tensor,
         *args,
         params: Optional["Packed"] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        th_ein: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        th_ein: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
diff --git a/src/caustics/lenses/tnfw.py b/src/caustics/lenses/tnfw.py
index 3c9478d6..cb613687 100644
--- a/src/caustics/lenses/tnfw.py
+++ b/src/caustics/lenses/tnfw.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="operator,union-attr"
 from math import pi
 from typing import Optional, Union
 
@@ -95,7 +96,7 @@ def __init__(
         s: float = 0.0,
         interpret_m_total_mass: bool = True,
         use_case="batchable",
-        name: str = None,
+        name: Optional[str] = None,
     ):
         """
         Initialize an instance of the TNFW lens class.
@@ -150,12 +151,12 @@ def get_concentration(
         self,
         *args,
         params: Optional[Packed] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        mass: Tensor = None,
-        scale_radius: Tensor = None,
-        tau: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        mass: Optional[Tensor] = None,
+        scale_radius: Optional[Tensor] = None,
+        tau: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -194,12 +195,12 @@ def get_truncation_radius(
         self,
         *args,
         params: Optional[Packed] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        mass: Tensor = None,
-        scale_radius: Tensor = None,
-        tau: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        mass: Optional[Tensor] = None,
+        scale_radius: Optional[Tensor] = None,
+        tau: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -234,12 +235,12 @@ def get_M0(
         self,
         *args,
         params: Optional[Packed] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        mass: Tensor = None,
-        scale_radius: Tensor = None,
-        tau: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        mass: Optional[Tensor] = None,
+        scale_radius: Optional[Tensor] = None,
+        tau: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -280,12 +281,12 @@ def get_scale_density(
         self,
         *args,
         params: Optional[Packed] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        mass: Tensor = None,
-        scale_radius: Tensor = None,
-        tau: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        mass: Optional[Tensor] = None,
+        scale_radius: Optional[Tensor] = None,
+        tau: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -324,12 +325,12 @@ def convergence(
         z_s: Tensor,
         *args,
         params: Optional[Packed] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        mass: Tensor = None,
-        scale_radius: Tensor = None,
-        tau: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        mass: Optional[Tensor] = None,
+        scale_radius: Optional[Tensor] = None,
+        tau: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -384,12 +385,12 @@ def mass_enclosed_2d(
         z_s: Tensor,
         *args,
         params: Optional[Packed] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        mass: Tensor = None,
-        scale_radius: Tensor = None,
-        tau: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        mass: Optional[Tensor] = None,
+        scale_radius: Optional[Tensor] = None,
+        tau: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
@@ -437,12 +438,12 @@ def physical_deflection_angle(
         z_s: Tensor,
         *args,
         params: Optional[Packed] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        mass: Tensor = None,
-        scale_radius: Tensor = None,
-        tau: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        mass: Optional[Tensor] = None,
+        scale_radius: Optional[Tensor] = None,
+        tau: Optional[Tensor] = None,
         **kwargs,
     ) -> tuple[Tensor, Tensor]:
         """Compute the physical deflection angle (arcsec) for this lens at
@@ -493,12 +494,12 @@ def potential(
         z_s: Tensor,
         *args,
         params: Optional[Packed] = None,
-        z_l: Tensor = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        mass: Tensor = None,
-        scale_radius: Tensor = None,
-        tau: Tensor = None,
+        z_l: Optional[Tensor] = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        mass: Optional[Tensor] = None,
+        scale_radius: Optional[Tensor] = None,
+        tau: Optional[Tensor] = None,
         **kwargs,
     ) -> Tensor:
         """
diff --git a/src/caustics/light/pixelated.py b/src/caustics/light/pixelated.py
index c445c4d0..7ea08408 100644
--- a/src/caustics/light/pixelated.py
+++ b/src/caustics/light/pixelated.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="union-attr"
 from typing import Optional, Union
 
 from torch import Tensor
@@ -42,7 +43,7 @@ def __init__(
         y0: Optional[Union[Tensor, float]] = None,
         pixelscale: Optional[Union[Tensor, float]] = None,
         shape: Optional[tuple[int, ...]] = None,
-        name: str = None,
+        name: Optional[str] = None,
     ):
         """
         Constructs the `Pixelated` object with the given parameters.
@@ -83,10 +84,10 @@ def brightness(
         y,
         *args,
         params: Optional["Packed"] = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        image: Tensor = None,
-        pixelscale: Tensor = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        image: Optional[Tensor] = None,
+        pixelscale: Optional[Tensor] = None,
         **kwargs,
     ):
         """
diff --git a/src/caustics/light/sersic.py b/src/caustics/light/sersic.py
index 11bc5e05..a5b8067c 100644
--- a/src/caustics/light/sersic.py
+++ b/src/caustics/light/sersic.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="operator,union-attr"
 from typing import Optional, Union
 
 from torch import Tensor
@@ -53,7 +54,7 @@ def __init__(
         Ie: Optional[Union[Tensor, float]] = None,
         s: float = 0.0,
         use_lenstronomy_k=False,
-        name: str = None,
+        name: Optional[str] = None,
     ):
         """
         Constructs the `Sersic` object with the given parameters.
@@ -100,13 +101,13 @@ def brightness(
         y,
         *args,
         params: Optional["Packed"] = None,
-        x0: Tensor = None,
-        y0: Tensor = None,
-        q: Tensor = None,
-        phi: Tensor = None,
-        n: Tensor = None,
-        Re: Tensor = None,
-        Ie: Tensor = None,
+        x0: Optional[Tensor] = None,
+        y0: Optional[Tensor] = None,
+        q: Optional[Tensor] = None,
+        phi: Optional[Tensor] = None,
+        n: Optional[Tensor] = None,
+        Re: Optional[Tensor] = None,
+        Ie: Optional[Tensor] = None,
         **kwargs,
     ):
         """
diff --git a/src/caustics/parameter.py b/src/caustics/parameter.py
index 933de529..f86c874b 100644
--- a/src/caustics/parameter.py
+++ b/src/caustics/parameter.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="union-attr"
 from typing import Optional, Union
 
 import torch
diff --git a/src/caustics/parametrized.py b/src/caustics/parametrized.py
index 5f9fab8f..6f00ccc8 100644
--- a/src/caustics/parametrized.py
+++ b/src/caustics/parametrized.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="var-annotated,index,type-arg"
 from collections import OrderedDict
 from math import prod
 from typing import Optional, Union
@@ -59,7 +60,7 @@ class Parametrized:
         Number of static parameters.
     """
 
-    def __init__(self, name: str = None):
+    def __init__(self, name: Optional[str] = None):
         if name is None:
             name = self._default_name()
         check_valid_name(name)
diff --git a/src/caustics/utils.py b/src/caustics/utils.py
index 42c04d1b..e16775f0 100644
--- a/src/caustics/utils.py
+++ b/src/caustics/utils.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="misc"
 from math import pi
 from typing import Callable, Optional, Tuple, Union
 from functools import partial, lru_cache