From 60de2b0b6a464b0922c4bbaad5c5564c3d00cb4f Mon Sep 17 00:00:00 2001
From: Martin Schubert <mfschubert@gmail.com>
Date: Fri, 1 Dec 2023 07:34:06 -0800
Subject: [PATCH 1/2] Serialize fmmax types

---
 src/invrs_gym/challenges/base.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/invrs_gym/challenges/base.py b/src/invrs_gym/challenges/base.py
index 54f48d0..ed3edea 100644
--- a/src/invrs_gym/challenges/base.py
+++ b/src/invrs_gym/challenges/base.py
@@ -6,9 +6,10 @@
 import abc
 from typing import Any, Callable, Dict, Tuple
 
+import fmmax
 import jax
 import jax.numpy as jnp
-from totypes import types
+from totypes import json_utils, types
 
 AuxDict = Dict[str, Any]
 PyTree = Any
@@ -74,3 +75,8 @@ def distance_to_target(self, response: Any) -> jnp.ndarray:
     @abc.abstractmethod
     def metrics(self, response: Any, params: PyTree, aux: AuxDict) -> AuxDict:
         """Compute metrics for a component response and associated quantities."""
+
+
+# Several challenges use the `fmmax` simulator, and contain `fmmax` custom objects in
+# their responses. Ensure these are serializable by registering with totypes.
+json_utils.register_custom_type(fmmax.basis.Expansion)

From ee53363596441ae2d4a921b15297bf8745c61842 Mon Sep 17 00:00:00 2001
From: Martin Schubert <mfschubert@gmail.com>
Date: Fri, 1 Dec 2023 07:45:16 -0800
Subject: [PATCH 2/2] Version updated from v0.3.0 to v0.3.1

---
 .bumpversion.toml         | 2 +-
 README.md                 | 2 +-
 pyproject.toml            | 2 +-
 src/invrs_gym/__init__.py | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/.bumpversion.toml b/.bumpversion.toml
index 068ee0b..16eeedd 100644
--- a/.bumpversion.toml
+++ b/.bumpversion.toml
@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "v0.3.0"
+current_version = "v0.3.1"
 commit = true
 commit_args = "--no-verify"
 tag = true
diff --git a/README.md b/README.md
index 7e715e9..60eecd2 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,5 @@
 # invrs-gym - A collection of inverse design challenges
-`v0.3.0`
+`v0.3.1`
 
 ## Overview
 The `invrs_gym` package is an open-source gym containing a diverse set of photonic design challenges, which are relevant for a wide range of applications such as AR/VR, optical networking, LIDAR, and others.
diff --git a/pyproject.toml b/pyproject.toml
index f4a8384..52bb275 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 
 name = "invrs_gym"
-version = "v0.3.0"
+version = "v0.3.1"
 description = "A collection of inverse design challenges"
 keywords = ["topology", "optimization", "jax", "inverse design"]
 readme = "README.md"
diff --git a/src/invrs_gym/__init__.py b/src/invrs_gym/__init__.py
index 6f0340a..b82b44d 100644
--- a/src/invrs_gym/__init__.py
+++ b/src/invrs_gym/__init__.py
@@ -3,7 +3,7 @@
 Copyright (c) 2023 The INVRS-IO authors.
 """
 
-__version__ = "v0.3.0"
+__version__ = "v0.3.1"
 __author__ = "Martin F. Schubert <mfschubert@gmail.com>"
 
 from invrs_gym import challenges as challenges