714        def serialize(self) -> dict[str, Any]:
715            result: dict[str, Any] = {
716                _FORMAT_KEY: f"{self.__class__.__name__}(SerializableDataclass)"
717            }
718            # for each field in the class
719            for field in dataclasses.fields(self):  # type: ignore[arg-type]
720                # need it to be our special SerializableField
721                if not isinstance(field, SerializableField):
722                    raise NotSerializableFieldException(
723                        f"Field '{field.name}' on class {self.__class__.__module__}.{self.__class__.__name__} is not a `SerializableField`, "
724                        f"but a {type(field)} "
725                        "this state should be inaccessible, please report this bug!"
726                    )
727
728                # try to save it
729                if field.serialize:
730                    try:
731                        # get the val
732                        value = getattr(self, field.name)
733                        # if it is a serializable dataclass, serialize it
734                        if isinstance(value, SerializableDataclass):
735                            value = value.serialize()
736                        # if the value has a serialization function, use that
737                        if hasattr(value, "serialize") and callable(value.serialize):
738                            value = value.serialize()
739                        # if the field has a serialization function, use that
740                        # it would be nice to be able to override a class's `.serialize()`, but that could lead to some inconsistencies!
741                        elif field.serialization_fn:
742                            value = field.serialization_fn(value)
743
744                        # store the value in the result
745                        result[field.name] = value
746                    except Exception as e:
747                        raise FieldSerializationError(
748                            "\n".join(
749                                [
750                                    f"Error serializing field '{field.name}' on class {self.__class__.__module__}.{self.__class__.__name__}",
751                                    f"{field = }",
752                                    f"{value = }",
753                                    f"{self = }",
754                                ]
755                            )
756                        ) from e
757
758            # store each property if we can get it
759            for prop in self._properties_to_serialize:
760                if hasattr(cls, prop):
761                    value = getattr(self, prop)
762                    result[prop] = value
763                else:
764                    raise AttributeError(
765                        f"Cannot serialize property '{prop}' on class {self.__class__.__module__}.{self.__class__.__name__}"
766                        + f"but it is in {self._properties_to_serialize = }"
767                        + f"\n{self = }"
768                    )
769
770            return result

777        @classmethod  # type: ignore[misc]
778        def load(cls, data: dict[str, Any] | T) -> Type[T]:
779            # HACK: this is kind of ugly, but it fixes a lot of issues for when we do recursive loading with ZANJ
780            if isinstance(data, cls):
781                return data
782
783            assert isinstance(
784                data, typing.Mapping
785            ), f"When loading {cls.__name__ = } expected a Mapping, but got {type(data) = }:\n{data = }"
786
787            cls_type_hints: dict[str, Any] = get_cls_type_hints(cls)
788
789            # initialize dict for keeping what we will pass to the constructor
790            ctor_kwargs: dict[str, Any] = dict()
791
792            # iterate over the fields of the class
793            for field in dataclasses.fields(cls):
794                # check if the field is a SerializableField
795                assert isinstance(
796                    field, SerializableField
797                ), f"Field '{field.name}' on class {cls.__name__} is not a SerializableField, but a {type(field)}. this state should be inaccessible, please report this bug!\nhttps://github.com/mivanit/muutils/issues/new"
798
799                # check if the field is in the data and if it should be initialized
800                if (field.name in data) and field.init:
801                    # get the value, we will be processing it
802                    value: Any = data[field.name]
803
804                    # get the type hint for the field
805                    field_type_hint: Any = cls_type_hints.get(field.name, None)
806
807                    # we rely on the init of `SerializableField` to check that only one of `loading_fn` and `deserialize_fn` is set
808                    if field.deserialize_fn:
809                        # if it has a deserialization function, use that
810                        value = field.deserialize_fn(value)
811                    elif field.loading_fn:
812                        # if it has a loading function, use that
813                        value = field.loading_fn(data)
814                    elif (
815                        field_type_hint is not None
816                        and hasattr(field_type_hint, "load")
817                        and callable(field_type_hint.load)
818                    ):
819                        # if no loading function but has a type hint with a load method, use that
820                        if isinstance(value, dict):
821                            value = field_type_hint.load(value)
822                        else:
823                            raise FieldLoadingError(
824                                f"Cannot load value into {field_type_hint}, expected {type(value) = } to be a dict\n{value = }"
825                            )
826                    else:
827                        # assume no loading needs to happen, keep `value` as-is
828                        pass
829
830                    # store the value in the constructor kwargs
831                    ctor_kwargs[field.name] = value
832
833            # create a new instance of the class with the constructor kwargs
834            output: cls = cls(**ctor_kwargs)
835
836            # validate the types of the fields if needed
837            if on_typecheck_mismatch != ErrorMode.IGNORE:
838                fields_valid: dict[str, bool] = (
839                    SerializableDataclass__validate_fields_types__dict(
840                        output,
841                        on_typecheck_error=on_typecheck_error,
842                    )
843                )
844
845                # if there are any fields that are not valid, raise an error
846                if not all(fields_valid.values()):
847                    msg: str = (
848                        f"Type mismatch in fields of {cls.__name__}:\n"
849                        + "\n".join(
850                            [
851                                f"{k}:\texpected {cls_type_hints[k] = }, but got value {getattr(output, k) = }, {type(getattr(output, k)) = }"
852                                for k, v in fields_valid.items()
853                                if not v
854                            ]
855                        )
856                    )
857
858                    on_typecheck_mismatch.process(
859                        msg, except_cls=FieldTypeMismatchError
860                    )
861
862            # return the new instance
863            return output

283def SerializableDataclass__validate_fields_types(
284    self: SerializableDataclass,
285    on_typecheck_error: ErrorMode = _DEFAULT_ON_TYPECHECK_ERROR,
286) -> bool:
287    """validate the types of all the fields on a `SerializableDataclass`. calls `SerializableDataclass__validate_field_type` for each field"""
288    return all(
289        SerializableDataclass__validate_fields_types__dict(
290            self, on_typecheck_error=on_typecheck_error
291        ).values()
292    )

154	def get_batch(
155		self,
156		idxs: list[int] | None,
157	) -> Float[np.ndarray, "in/tgt=2 item x y rgb=3"]:
158		"""get a batch of mazes as a tensor, from a list of indices"""
159		if idxs is None:
160			idxs = list(range(len(self)))
161
162		inputs: list[Float[np.ndarray, "x y rgb=3"]]
163		targets: list[Float[np.ndarray, "x y rgb=3"]]
164		inputs, targets = zip(*[self[i] for i in idxs], strict=False)  # type: ignore[assignment]
165
166		return np.array([inputs, targets])

169	@classmethod
170	def from_config(
171		cls,
172		cfg: RasterizedMazeDatasetConfig | MazeDatasetConfig,  # type: ignore[override]
173		do_generate: bool = True,
174		load_local: bool = True,
175		save_local: bool = True,
176		zanj: ZANJ | None = None,
177		do_download: bool = True,
178		local_base_path: Path = Path("data/maze_dataset"),
179		except_on_config_mismatch: bool = True,
180		allow_generation_metadata_filter_mismatch: bool = True,
181		verbose: bool = False,
182		**kwargs,
183	) -> "RasterizedMazeDataset":
184		"""create a rasterized maze dataset from a config
185
186		priority of loading:
187		1. load from local
188		2. download
189		3. generate
190
191		"""
192		return typing.cast(
193			RasterizedMazeDataset,
194			super().from_config(
195				cfg=cfg,
196				do_generate=do_generate,
197				load_local=load_local,
198				save_local=save_local,
199				zanj=zanj,
200				do_download=do_download,
201				local_base_path=local_base_path,
202				except_on_config_mismatch=except_on_config_mismatch,
203				allow_generation_metadata_filter_mismatch=allow_generation_metadata_filter_mismatch,
204				verbose=verbose,
205				**kwargs,
206			),
207		)

209	@classmethod
210	def from_config_augmented(
211		cls,
212		cfg: RasterizedMazeDatasetConfig,
213		**kwargs,
214	) -> "RasterizedMazeDataset":
215		"""loads either a maze transformer dataset or an easy_2_hard dataset"""
216		_cfg_temp: MazeDatasetConfig = MazeDatasetConfig.load(cfg.serialize())
217		return cls.from_base_MazeDataset(
218			cls.from_config(cfg=_cfg_temp, **kwargs),
219			added_params={
220				k: v
221				for k, v in cfg.serialize().items()
222				if k in _RASTERIZED_CFG_ADDED_PARAMS
223			},
224		)

226	@classmethod
227	def from_base_MazeDataset(
228		cls,
229		base_dataset: MazeDataset,
230		added_params: dict | None = None,
231	) -> "RasterizedMazeDataset":
232		"""loads either a maze transformer dataset or an easy_2_hard dataset"""
233		if added_params is None:
234			added_params = dict(
235				remove_isolated_cells=True,
236				extend_pixels=True,
237			)
238		cfg: RasterizedMazeDatasetConfig = RasterizedMazeDatasetConfig.load(
239			{
240				**base_dataset.cfg.serialize(),
241				**added_params,
242			},
243		)
244		output: RasterizedMazeDataset = cls(
245			cfg=cfg,
246			mazes=base_dataset.mazes,
247		)
248		return output

250	def plot(self, count: int | None = None, show: bool = True) -> tuple | None:
251		"""plot the first `count` mazes in the dataset"""
252		import matplotlib.pyplot as plt
253
254		print(f"{self[0][0].shape = }, {self[0][1].shape = }")
255		count = count or len(self)
256		if count == 0:
257			print("No mazes to plot for dataset")
258			return None
259		fig, axes = plt.subplots(2, count, figsize=(15, 5))
260		if count == 1:
261			axes = [axes]
262		for i in range(count):
263			axes[0, i].imshow(self[i][0])
264			axes[1, i].imshow(self[i][1])
265			# remove ticks
266			axes[0, i].set_xticks([])
267			axes[0, i].set_yticks([])
268			axes[1, i].set_xticks([])
269			axes[1, i].set_yticks([])
270
271		if show:
272			plt.show()
273
274		return fig, axes