How can I prevent infinite re-renders for my Jotai Map of Atoms extraction hook?

Question

I am using Jotai state management with my React app.

An overview of the pattern. I have a range of form components that consume Atoms through their props. These form components then convert these atoms into something usable through

const [formInput,setFormInput] = useAtom(formAtom)

These forms can be dynamically generated and there could be multiple instances of the same form so Atoms must be dynamically generated. What I have set up is a Map of atoms

atom<Map<string, PrimitiveAtom<Record<string, unknown>>>>(new Map())

I then use a custom hook to attach new atoms onto this map or to return the atom if it already exists within the map.

import { atom, PrimitiveAtom, useAtom } from 'jotai';
import { useCallback, useEffect, useMemo, useState } from 'react';

const useSetGetAddAtom = (
    atomMapOfAtoms: PrimitiveAtom<Map<string, PrimitiveAtom<Record<string, unknown>>>>,
    formId: string
): {
    formAtom: PrimitiveAtom<Record<string, unknown>>;
    generateFormAtom: () => PrimitiveAtom<Record<string, unknown>>;
} => {
    const [formAtom, setFormAtom] = useState<PrimitiveAtom<Record<string, unknown>> | null>(null);
    const [mapOfAtoms, setMapOfAtoms] = useAtom<Map<string, PrimitiveAtom<Record<string, unknown>>>>(atomMapOfAtoms);

    const generateFormAtom = useCallback(() => {
        if (!mapOfAtoms.has(formId)) {
            // I have tried initializing this new atom at the top of the hook however it has not made a difference
            const newAtom = atom<Record<string, unknown>>({});
            setMapOfAtoms((prev) => {
                const newMap = new Map(prev);
                newMap.set(formId, newAtom);
                return newMap;
            });
            return newAtom;
        } else {
            return mapOfAtoms.get(formId)!;
        }
    }, [mapOfAtoms, setMapOfAtoms, formId]);

    // It wasn't happy
    useEffect(() => {
        const atom = generateFormAtom();
        setFormAtom(atom);
    }, [generateFormAtom]);

    // Render a placeholder or loading state until formAtom is set - this itself doesn't cause the infinite rerenders as it will only be blank on the very initial render. Without the dummy atom everything breaks due to the hook returning null.

    if (!formAtom) {
        return { formAtom: atom({}), generateFormAtom };
    }

    return { formAtom, generateFormAtom };
};

export default useSetGetAddAtom;

So when I use this, in specific circumstances it will cause an infinite re-render of a function. I have a suspicion that when a parent and child both are consuming the same piece of state, that the switch over from the dummy atom to the atom in the map causes a render in the child, which might cause the parent to re-render or something along those lines.

I return the actual atom and not the values as the forms consume the atoms and convert them themselves.

This feels very complicated just to have combined state that can communicate together and I wonder have I run myself into some kind of Anti-pattern. If there are better ways of managing atoms together I would be very interested as I can't seem to find much documentation on it.

I have tried memoizing any atoms that are composed within the render function as well as this is the most common answer however the infinite re-renders seem linked to the hook.

I have recently also tried simplifying the hook, however this has not solved the re-render issue yet.

import { useEffect, useMemo } from 'react';
import { atom, useAtom } from 'jotai';

import { FormAtomMap } from '@/types/formTypes.ts';

const useCreateOrFetchAtom = (atomMap: FormAtomMap, atomId: string) => {
  const [formMap, setFormMap] = useAtom(atomMap);
  const newAtom = useMemo(() => atom<Record<string, unknown>>({}), []);

  useEffect(() => {
    setFormMap((prev) => {
      if (prev.has(atomId)) {
        return prev;
      } else {
        const newMap = new Map();
        newMap.set(atomId, newAtom);
        return newMap;
      }
    });
  }, [newAtom, atomId, setFormMap]);

  return formMap.get(atomId) || newAtom;
};

export default useCreateOrFetchAtom;