diff --git a/.changeset/fluffy-llamas-dance.md b/.changeset/fluffy-llamas-dance.md
new file mode 100644
index 0000000..13b6b45
--- /dev/null
+++ b/.changeset/fluffy-llamas-dance.md
@@ -0,0 +1,13 @@
+---
+'jotai-x': minor
+---
+
+- Atoms other than `atom` can now be passed in the `initialState` argument to `createAtomStore`. Primitive values use `atom` by default
+- Added an `extend` option to `createAtomStore` that lets you add derived atoms to the store
+- New accessors on `UseStoreApi`
+  - `useMyStore().store()` returns the `JotaiStore` for the current context, or undefined if no store exists
+  - `useMyStore().{get,set,use}.atom(someAtom)` accesses `someAtom` through the store
+- Types: remove exports for some internal types
+  - `GetRecord`
+  - `SetRecord`
+  - `UseRecord`
diff --git a/README.md b/README.md
index 6079108..5f0a691 100644
--- a/README.md
+++ b/README.md
@@ -54,26 +54,25 @@ createAtomStore<T extends object>(initialState: T, options?: CreateAtomStoreOpti
 The **`options`** object can include several properties to customize the behavior of your store:
 
 - **`name`**: A string representing the name of the store, which can be helpful for debugging or when working with multiple stores.
-- **`store`**: Allows specifying a [Jotai store](https://jotai.org/docs/core/store) if you want to use a custom one. Optional.
 - **`delay`**: If you need to introduce a delay in state updates, you can specify it here. Optional.
 - **`effect`**: A React component that can be used to run effects inside the provider. Optional.
+- **`extend`**: Extend the store with derived atoms based on the store state. Optional.
 
 #### Return Value
 
 The **`createAtomStore`** function returns an object (**`AtomStoreApi`**) containing the following properties and methods for interacting with the store:
 
 - **`use<Name>Store`**: 
-  - A function that returns the following objects: **`get`**, **`set`**, and **`use`**, where values are hooks for each state defined in the store. 
+  - A function that returns the following objects: **`get`**, **`set`**, **`use`** and **`store`**, where values are hooks for each state defined in the store.
   - **`get`**: Hooks for accessing a state within a component,  ensuring re-rendering when the state changes. See [useAtomValue](https://jotai.org/docs/core/use-atom#useatomvalue).
   - **`set`**: Hooks for setting a state within a component. See [useSetAtom](https://jotai.org/docs/core/use-atom#usesetatom).
   - **`use`**: Hooks for accessing and setting a state within a component, ensuring re-rendering when the state changes. See [useAtom](https://jotai.org/docs/core/use-atom).
+  - **`store`**: A hook to access the [JotaiStore](https://jotai.org/docs/core/store) for the current context.
   - Example: `const [element, setElement] = useElementStore().use.element()`
 - **`<Name>Provider`**:
   - The API includes dynamically generated provider components for each defined store. This allows  scoped state management within your application. More information in the next section.
 - **`<name>Store`**:
-  - Advanced API you generally don't need.
-  - **`atom`**: A hook for accessing state within a component,  ensuring re-rendering when the state changes. See [atom](https://jotai.org/docs/core/atom).
-  - **`extend`**: Extends the store with additional atoms. 
+  - **`atom`**: Access the atoms used by the store, including derived atoms defined using `extend`. See [atom](https://jotai.org/docs/core/atom).
   
 ### **Provider-Based Store Hydration and Synchronization**
 
@@ -86,6 +85,40 @@ The **`createAtomStore`** function returns an object (**`AtomStoreApi`**) contai
 
 JotaiX creates scoped providers, enabling more granular control over different segments of state within your application. `createAtomStore` sets up a context for each store, which can be scoped using the **`scope`** prop. This is particularly beneficial in complex applications where nested providers are needed.
 
+### Derived Atoms
+
+There are two ways of creating derived atoms from your JotaiX store.
+
+#### Derived Atoms Using `extend`
+
+Atoms defined using the `extend` option are made available in the same places as other values in the store.
+
+```ts
+const { useUserStore } = createAtomStore({
+  username: 'Alice',
+}, {
+  name: 'user',
+  extend: (atoms) => ({
+    intro: atom((get) => `My name is ${get(atoms.username)}`),
+  }),
+});
+
+const intro = useAppStore().get.intro();
+```
+
+#### Externally Defined Derived Atoms
+
+Derived atoms can also be defined externally by accessing the store's atoms through the `<name>Store` API. Externally defined atoms can be accessed through the store using the special `use<Name>Store().{get,set,use}.atom` hooks.
+
+```ts
+const { userStore, useUserStore } = createAtomStore({
+  username: 'Alice',
+}, { name: 'user' });
+
+const introAtom = atom((get) => `My name is ${get(userStore.atom.username)}`);
+const intro = useUserStore().get.atom(introAtom);
+```
+
 ### Example Usage
 
 #### 1. Create a store
@@ -195,11 +228,6 @@ const Component = () => {
 
 ## Contributing
 
-### Roadmap
-
-- [ ] Support other atoms like `atomWithStorage`
-- [ ] Improve `extend` API to be more modular.
-
 ### Ideas and discussions
 
 [Discussions](https://github.com/udecode/jotai-x/discussions) is the best
diff --git a/package.json b/package.json
index 76d710a..5754635 100644
--- a/package.json
+++ b/package.json
@@ -18,6 +18,7 @@
     "typecheck": "yarn g:typecheck",
     "typedoc": "npx typedoc --options scripts/typedoc.json",
     "release": "yarn build && yarn changeset publish",
+    "release:next": "yarn build && yarn changeset publish --tag next",
     "g:brl": "turbo --filter \"./packages/**\" brl --no-daemon",
     "g:build": "turbo --filter \"./packages/**\" build --no-daemon",
     "g:build:watch": "yarn build:watch",
@@ -25,7 +26,6 @@
     "g:clean": "yarn clean:turbo && turbo --filter \"./packages/**\" clean --no-daemon",
     "g:lint": "turbo --filter \"./packages/**\" lint --no-daemon",
     "g:lint:fix": "turbo lint:fix --no-daemon",
-    "g:release:next": "yarn yarn build && yarn changeset publish --tag next",
     "g:test": "turbo --filter \"./packages/**\" test --no-daemon",
     "g:test:watch": "turbo --filter \"./packages/**\" test:watch --no-daemon",
     "g:test:cov": "yarn g:test --coverage",
diff --git a/packages/jotai-x/README.md b/packages/jotai-x/README.md
index 6079108..5f0a691 100644
--- a/packages/jotai-x/README.md
+++ b/packages/jotai-x/README.md
@@ -54,26 +54,25 @@ createAtomStore<T extends object>(initialState: T, options?: CreateAtomStoreOpti
 The **`options`** object can include several properties to customize the behavior of your store:
 
 - **`name`**: A string representing the name of the store, which can be helpful for debugging or when working with multiple stores.
-- **`store`**: Allows specifying a [Jotai store](https://jotai.org/docs/core/store) if you want to use a custom one. Optional.
 - **`delay`**: If you need to introduce a delay in state updates, you can specify it here. Optional.
 - **`effect`**: A React component that can be used to run effects inside the provider. Optional.
+- **`extend`**: Extend the store with derived atoms based on the store state. Optional.
 
 #### Return Value
 
 The **`createAtomStore`** function returns an object (**`AtomStoreApi`**) containing the following properties and methods for interacting with the store:
 
 - **`use<Name>Store`**: 
-  - A function that returns the following objects: **`get`**, **`set`**, and **`use`**, where values are hooks for each state defined in the store. 
+  - A function that returns the following objects: **`get`**, **`set`**, **`use`** and **`store`**, where values are hooks for each state defined in the store.
   - **`get`**: Hooks for accessing a state within a component,  ensuring re-rendering when the state changes. See [useAtomValue](https://jotai.org/docs/core/use-atom#useatomvalue).
   - **`set`**: Hooks for setting a state within a component. See [useSetAtom](https://jotai.org/docs/core/use-atom#usesetatom).
   - **`use`**: Hooks for accessing and setting a state within a component, ensuring re-rendering when the state changes. See [useAtom](https://jotai.org/docs/core/use-atom).
+  - **`store`**: A hook to access the [JotaiStore](https://jotai.org/docs/core/store) for the current context.
   - Example: `const [element, setElement] = useElementStore().use.element()`
 - **`<Name>Provider`**:
   - The API includes dynamically generated provider components for each defined store. This allows  scoped state management within your application. More information in the next section.
 - **`<name>Store`**:
-  - Advanced API you generally don't need.
-  - **`atom`**: A hook for accessing state within a component,  ensuring re-rendering when the state changes. See [atom](https://jotai.org/docs/core/atom).
-  - **`extend`**: Extends the store with additional atoms. 
+  - **`atom`**: Access the atoms used by the store, including derived atoms defined using `extend`. See [atom](https://jotai.org/docs/core/atom).
   
 ### **Provider-Based Store Hydration and Synchronization**
 
@@ -86,6 +85,40 @@ The **`createAtomStore`** function returns an object (**`AtomStoreApi`**) contai
 
 JotaiX creates scoped providers, enabling more granular control over different segments of state within your application. `createAtomStore` sets up a context for each store, which can be scoped using the **`scope`** prop. This is particularly beneficial in complex applications where nested providers are needed.
 
+### Derived Atoms
+
+There are two ways of creating derived atoms from your JotaiX store.
+
+#### Derived Atoms Using `extend`
+
+Atoms defined using the `extend` option are made available in the same places as other values in the store.
+
+```ts
+const { useUserStore } = createAtomStore({
+  username: 'Alice',
+}, {
+  name: 'user',
+  extend: (atoms) => ({
+    intro: atom((get) => `My name is ${get(atoms.username)}`),
+  }),
+});
+
+const intro = useAppStore().get.intro();
+```
+
+#### Externally Defined Derived Atoms
+
+Derived atoms can also be defined externally by accessing the store's atoms through the `<name>Store` API. Externally defined atoms can be accessed through the store using the special `use<Name>Store().{get,set,use}.atom` hooks.
+
+```ts
+const { userStore, useUserStore } = createAtomStore({
+  username: 'Alice',
+}, { name: 'user' });
+
+const introAtom = atom((get) => `My name is ${get(userStore.atom.username)}`);
+const intro = useUserStore().get.atom(introAtom);
+```
+
 ### Example Usage
 
 #### 1. Create a store
@@ -195,11 +228,6 @@ const Component = () => {
 
 ## Contributing
 
-### Roadmap
-
-- [ ] Support other atoms like `atomWithStorage`
-- [ ] Improve `extend` API to be more modular.
-
 ### Ideas and discussions
 
 [Discussions](https://github.com/udecode/jotai-x/discussions) is the best
diff --git a/packages/jotai-x/src/atomWithFn.ts b/packages/jotai-x/src/atomWithFn.ts
new file mode 100644
index 0000000..5c759a1
--- /dev/null
+++ b/packages/jotai-x/src/atomWithFn.ts
@@ -0,0 +1,32 @@
+import { atom } from 'jotai';
+
+import type { WritableAtom } from 'jotai/vanilla';
+
+type WrapFn<T> = T extends (...args: infer _A) => infer _R ? { __fn: T } : T;
+
+const wrapFn = <T>(fnOrValue: T): WrapFn<T> =>
+  (typeof fnOrValue === 'function' ? { __fn: fnOrValue } : fnOrValue) as any;
+
+type UnwrapFn<T> = T extends { __fn: infer U } ? U : T;
+
+const unwrapFn = <T>(wrappedFnOrValue: T): UnwrapFn<T> =>
+  (wrappedFnOrValue &&
+  typeof wrappedFnOrValue === 'object' &&
+  '__fn' in wrappedFnOrValue
+    ? wrappedFnOrValue.__fn
+    : wrappedFnOrValue) as any;
+
+/**
+ * Jotai atoms don't allow functions as values by default. This function is a
+ * drop-in replacement for `atom` that wraps functions in an object while
+ * leaving non-functions unchanged. The wrapper object should be completely
+ * invisible to consumers of the atom.
+ */
+export const atomWithFn = <T>(initialValue: T): WritableAtom<T, [T], void> => {
+  const baseAtom = atom(wrapFn(initialValue));
+
+  return atom(
+    (get) => unwrapFn(get(baseAtom)) as T,
+    (_get, set, value) => set(baseAtom, wrapFn(value))
+  );
+};
diff --git a/packages/jotai-x/src/createAtomProvider.tsx b/packages/jotai-x/src/createAtomProvider.tsx
index 5412023..efcc6d8 100644
--- a/packages/jotai-x/src/createAtomProvider.tsx
+++ b/packages/jotai-x/src/createAtomProvider.tsx
@@ -9,7 +9,7 @@ import React, {
 import { createStore } from 'jotai/vanilla';
 
 import { AtomProvider, AtomProviderProps } from './atomProvider';
-import { AtomRecord, JotaiStore } from './createAtomStore';
+import { JotaiStore, SimpleWritableAtomRecord } from './createAtomStore';
 import { useHydrateStore, useSyncStore } from './useHydrateStore';
 
 const getFullyQualifiedScope = (storeName: string, scope: string) => {
@@ -62,7 +62,7 @@ export const HydrateAtoms = <T extends object>({
   atoms,
   ...props
 }: Omit<ProviderProps<T>, 'scope'> & {
-  atoms: AtomRecord<T>;
+  atoms: SimpleWritableAtomRecord<T>;
 }) => {
   useHydrateStore(atoms, { ...initialValues, ...props } as any, {
     store,
@@ -81,7 +81,7 @@ export const HydrateAtoms = <T extends object>({
  */
 export const createAtomProvider = <T extends object, N extends string = ''>(
   storeScope: N,
-  atoms: AtomRecord<T>,
+  atoms: SimpleWritableAtomRecord<T>,
   options: { effect?: FC } = {}
 ) => {
   const Effect = options.effect;
diff --git a/packages/jotai-x/src/createAtomStore.spec.tsx b/packages/jotai-x/src/createAtomStore.spec.tsx
index 9d911f2..b519f32 100644
--- a/packages/jotai-x/src/createAtomStore.spec.tsx
+++ b/packages/jotai-x/src/createAtomStore.spec.tsx
@@ -1,7 +1,9 @@
 import '@testing-library/jest-dom';
 
 import React, { ReactNode, useState } from 'react';
-import { act, render, renderHook } from '@testing-library/react';
+import { act, queryByText, render, renderHook } from '@testing-library/react';
+import { atom, PrimitiveAtom, useAtomValue } from 'jotai';
+import { splitAtom } from 'jotai/utils';
 
 import { createAtomStore } from './createAtomStore';
 
@@ -10,6 +12,7 @@ describe('createAtomStore', () => {
     type MyTestStoreValue = {
       name: string;
       age: number;
+      becomeFriends: () => void;
     };
 
     const INITIAL_NAME = 'John';
@@ -18,16 +21,15 @@ describe('createAtomStore', () => {
     const initialTestStoreValue: MyTestStoreValue = {
       name: INITIAL_NAME,
       age: INITIAL_AGE,
+      becomeFriends: () => {},
     };
 
-    const { useMyTestStoreStore, MyTestStoreProvider } = createAtomStore(
-      initialTestStoreValue,
-      { name: 'myTestStore' as const }
-    );
+    const { myTestStoreStore, useMyTestStoreStore, MyTestStoreProvider } =
+      createAtomStore(initialTestStoreValue, { name: 'myTestStore' as const });
 
     const ReadOnlyConsumer = () => {
-      const [name] = useMyTestStoreStore().use.name();
-      const [age] = useMyTestStoreStore().use.age();
+      const name = useMyTestStoreStore().get.name();
+      const age = useMyTestStoreStore().get.age();
 
       return (
         <div>
@@ -69,6 +71,76 @@ describe('createAtomStore', () => {
       );
     };
 
+    const BecomeFriendsProvider = ({ children }: { children: ReactNode }) => {
+      const [becameFriends, setBecameFriends] = useState(false);
+
+      return (
+        <>
+          <MyTestStoreProvider becomeFriends={() => setBecameFriends(true)}>
+            {children}
+          </MyTestStoreProvider>
+
+          <div>becameFriends: {becameFriends.toString()}</div>
+        </>
+      );
+    };
+
+    const BecomeFriendsGetter = () => {
+      // Make sure both of these are actual functions, not wrapped functions
+      const becomeFriends1 = useMyTestStoreStore().get.becomeFriends();
+      const becomeFriends2 = useMyTestStoreStore().get.atom(
+        myTestStoreStore.atom.becomeFriends
+      );
+
+      return (
+        <button
+          type="button"
+          onClick={() => {
+            becomeFriends1();
+            becomeFriends2();
+          }}
+        >
+          Become Friends
+        </button>
+      );
+    };
+
+    const BecomeFriendsSetter = () => {
+      const setBecomeFriends = useMyTestStoreStore().set.becomeFriends();
+      const [becameFriends, setBecameFriends] = useState(false);
+
+      return (
+        <>
+          <button
+            type="button"
+            onClick={() => setBecomeFriends(() => setBecameFriends(true))}
+          >
+            Change Callback
+          </button>
+
+          <div>setterBecameFriends: {becameFriends.toString()}</div>
+        </>
+      );
+    };
+
+    const BecomeFriendsUser = () => {
+      const [, setBecomeFriends] = useMyTestStoreStore().use.becomeFriends();
+      const [becameFriends, setBecameFriends] = useState(false);
+
+      return (
+        <>
+          <button
+            type="button"
+            onClick={() => setBecomeFriends(() => setBecameFriends(true))}
+          >
+            Change Callback
+          </button>
+
+          <div>userBecameFriends: {becameFriends.toString()}</div>
+        </>
+      );
+    };
+
     beforeEach(() => {
       renderHook(() => useMyTestStoreStore().set.name()(INITIAL_NAME));
       renderHook(() => useMyTestStoreStore().set.age()(INITIAL_AGE));
@@ -155,6 +227,46 @@ describe('createAtomStore', () => {
       expect(getByText(INITIAL_NAME)).toBeInTheDocument();
       expect(getByText(WRITE_ONLY_CONSUMER_AGE)).toBeInTheDocument();
     });
+
+    it('provides and gets functions', () => {
+      const { getByText } = render(
+        <BecomeFriendsProvider>
+          <BecomeFriendsGetter />
+        </BecomeFriendsProvider>
+      );
+
+      expect(getByText('becameFriends: false')).toBeInTheDocument();
+      act(() => getByText('Become Friends').click());
+      expect(getByText('becameFriends: true')).toBeInTheDocument();
+    });
+
+    it('sets functions', () => {
+      const { getByText } = render(
+        <BecomeFriendsProvider>
+          <BecomeFriendsSetter />
+          <BecomeFriendsGetter />
+        </BecomeFriendsProvider>
+      );
+
+      act(() => getByText('Change Callback').click());
+      expect(getByText('setterBecameFriends: false')).toBeInTheDocument();
+      act(() => getByText('Become Friends').click());
+      expect(getByText('setterBecameFriends: true')).toBeInTheDocument();
+    });
+
+    it('uses functions', () => {
+      const { getByText } = render(
+        <BecomeFriendsProvider>
+          <BecomeFriendsUser />
+          <BecomeFriendsGetter />
+        </BecomeFriendsProvider>
+      );
+
+      act(() => getByText('Change Callback').click());
+      expect(getByText('userBecameFriends: false')).toBeInTheDocument();
+      act(() => getByText('Become Friends').click());
+      expect(getByText('userBecameFriends: true')).toBeInTheDocument();
+    });
   });
 
   describe('scoped providers', () => {
@@ -170,7 +282,7 @@ describe('createAtomStore', () => {
       });
 
     const ReadOnlyConsumer = ({ scope }: { scope: string }) => {
-      const [age] = useMyScopedTestStoreStore().use.age({ scope });
+      const age = useMyScopedTestStoreStore().get.age({ scope });
 
       return (
         <div>
@@ -184,7 +296,7 @@ describe('createAtomStore', () => {
     }: {
       scope: string;
     }) => {
-      const [age] = useMyScopedTestStoreStore(scope).use.age();
+      const age = useMyScopedTestStoreStore(scope).get.age();
 
       return (
         <div>
@@ -269,7 +381,7 @@ describe('createAtomStore', () => {
       });
 
     const FirstReadOnlyConsumer = () => {
-      const [name] = useMyFirstTestStoreStore().use.name();
+      const name = useMyFirstTestStoreStore().get.name();
 
       return (
         <div>
@@ -279,7 +391,7 @@ describe('createAtomStore', () => {
     };
 
     const SecondReadOnlyConsumer = () => {
-      const [age] = useMySecondTestStoreStore().use.age();
+      const age = useMySecondTestStoreStore().get.age();
 
       return (
         <div>
@@ -302,4 +414,299 @@ describe('createAtomStore', () => {
       expect(getByText('98')).toBeInTheDocument();
     });
   });
+
+  describe('extended stores', () => {
+    type User = {
+      name: string;
+      age: number;
+    };
+
+    const initialUser: User = {
+      name: 'Jane',
+      age: 98,
+    };
+
+    const { userStore, useUserStore, UserProvider } = createAtomStore(
+      initialUser,
+      {
+        name: 'user' as const,
+        extend: ({ name, age }) => ({
+          bio: atom((get) => `${get(name)} is ${get(age)} years old`),
+        }),
+      }
+    );
+
+    const ReadOnlyConsumer = () => {
+      const bio = useUserStore().get.bio();
+
+      return <div>{bio}</div>;
+    };
+
+    it('includes extended atom in store object', () => {
+      const { result } = renderHook(() => useAtomValue(userStore.atom.bio));
+      expect(result.current).toBe('Jane is 98 years old');
+    });
+
+    it('includes extended atom in get hooks', () => {
+      const { result } = renderHook(() => useUserStore().get.bio());
+      expect(result.current).toBe('Jane is 98 years old');
+    });
+
+    it('does not include read-only extended atom in set hooks', () => {
+      const { result } = renderHook(() => Object.keys(useUserStore().set));
+      expect(result.current).not.toContain('bio');
+    });
+
+    it('does not include read-only extended atom in use hooks', () => {
+      const { result } = renderHook(() => Object.keys(useUserStore().use));
+      expect(result.current).not.toContain('bio');
+    });
+
+    it('computes extended atom based on current state', () => {
+      const { getByText } = render(
+        <UserProvider name="John" age={42}>
+          <ReadOnlyConsumer />
+        </UserProvider>
+      );
+
+      expect(getByText('John is 42 years old')).toBeInTheDocument();
+    });
+  });
+
+  describe('passing atoms as part of initial state', () => {
+    type CustomAtom<T> = PrimitiveAtom<T> & {
+      isCustomAtom: true;
+    };
+
+    const createCustomAtom = <T,>(value: T): CustomAtom<T> => ({
+      ...atom(value),
+      isCustomAtom: true,
+    });
+
+    const { customStore } = createAtomStore(
+      {
+        x: createCustomAtom(1),
+      },
+      {
+        name: 'custom' as const,
+      }
+    );
+
+    it('uses passed atom', () => {
+      const myAtom = customStore.atom.x as CustomAtom<number>;
+      expect(myAtom.isCustomAtom).toBe(true);
+    });
+  });
+
+  describe('arbitrary atom accessors', () => {
+    type User = {
+      name: string;
+    };
+
+    const initialUser: User = {
+      name: 'Jane',
+    };
+
+    const { userStore, useUserStore, UserProvider } = createAtomStore(
+      initialUser,
+      {
+        name: 'user' as const,
+      }
+    );
+
+    const derivedAtom = atom((get) => `My name is ${get(userStore.atom.name)}`);
+
+    const DerivedAtomConsumer = () => {
+      const message = useUserStore().get.atom(derivedAtom);
+
+      return <div>{message}</div>;
+    };
+
+    it('accesses arbitrary atom within store', () => {
+      const { getByText } = render(
+        <UserProvider name="John">
+          <DerivedAtomConsumer />
+        </UserProvider>
+      );
+
+      expect(getByText('My name is John')).toBeInTheDocument();
+    });
+  });
+
+  describe('splitAtoms using todoStore.atom.items', () => {
+    const initialState = {
+      items: [] as {
+        task: string;
+        done: boolean;
+      }[],
+    };
+
+    const { todoStore, useTodoStore, TodoProvider } = createAtomStore(
+      initialState,
+      {
+        name: 'todo' as const,
+      }
+    );
+
+    const todoAtomsAtom = splitAtom(todoStore.atom.items);
+
+    type TodoType = (typeof initialState)['items'][number];
+
+    const TodoItem = ({
+      todoAtom,
+      remove,
+    }: {
+      todoAtom: PrimitiveAtom<TodoType>;
+      remove: () => void;
+    }) => {
+      const [todo, setTodo] = useTodoStore().use.atom(todoAtom);
+
+      return (
+        <div>
+          <label>{todo.task}</label>
+          <input
+            type="checkbox"
+            checked={todo.done}
+            onChange={() => {
+              setTodo((oldValue) => ({ ...oldValue, done: !oldValue.done }));
+            }}
+          />
+          {/* eslint-disable-next-line react/button-has-type */}
+          <button onClick={remove}>remove {todo.task}</button>
+        </div>
+      );
+    };
+
+    const TodoList = () => {
+      const [todoAtoms, dispatch] = useTodoStore().use.atom(todoAtomsAtom);
+      return (
+        <ul>
+          {todoAtoms.map((todoAtom) => (
+            <TodoItem
+              key={`${todoAtom}`}
+              todoAtom={todoAtom}
+              remove={() => dispatch({ type: 'remove', atom: todoAtom })}
+            />
+          ))}
+        </ul>
+      );
+    };
+
+    it('should work', () => {
+      const { getByText, container } = render(
+        <TodoProvider
+          initialValues={{
+            items: [
+              {
+                task: 'help the town',
+                done: false,
+              },
+              {
+                task: 'feed the dragon',
+                done: false,
+              },
+            ],
+          }}
+        >
+          <TodoList />
+        </TodoProvider>
+      );
+
+      expect(getByText('help the town')).toBeInTheDocument();
+      expect(getByText('feed the dragon')).toBeInTheDocument();
+
+      act(() => getByText('remove help the town').click());
+
+      expect(queryByText(container, 'help the town')).not.toBeInTheDocument();
+      expect(getByText('feed the dragon')).toBeInTheDocument();
+    });
+  });
+
+  describe('splitAtoms using extend', () => {
+    const initialState = {
+      items: [] as {
+        task: string;
+        done: boolean;
+      }[],
+    };
+
+    const { useTodoStore, TodoProvider } = createAtomStore(initialState, {
+      name: 'todo' as const,
+      extend: ({ items }) => ({
+        itemAtoms: splitAtom(items),
+      }),
+    });
+
+    type TodoType = (typeof initialState)['items'][number];
+
+    const TodoItem = ({
+      todoAtom,
+      remove,
+    }: {
+      todoAtom: PrimitiveAtom<TodoType>;
+      remove: () => void;
+    }) => {
+      const [todo, setTodo] = useTodoStore().use.atom(todoAtom);
+
+      return (
+        <div>
+          <label>{todo.task}</label>
+          <input
+            type="checkbox"
+            checked={todo.done}
+            onChange={() => {
+              setTodo((oldValue) => ({ ...oldValue, done: !oldValue.done }));
+            }}
+          />
+          {/* eslint-disable-next-line react/button-has-type */}
+          <button onClick={remove}>remove {todo.task}</button>
+        </div>
+      );
+    };
+
+    const TodoList = () => {
+      const [todoAtoms, dispatch] = useTodoStore().use.itemAtoms();
+
+      return (
+        <ul>
+          {todoAtoms.map((todoAtom) => (
+            <TodoItem
+              key={`${todoAtom}`}
+              todoAtom={todoAtom}
+              remove={() => dispatch({ type: 'remove', atom: todoAtom })}
+            />
+          ))}
+        </ul>
+      );
+    };
+
+    it('should work', () => {
+      const { getByText, container } = render(
+        <TodoProvider
+          initialValues={{
+            items: [
+              {
+                task: 'help the town',
+                done: false,
+              },
+              {
+                task: 'feed the dragon',
+                done: false,
+              },
+            ],
+          }}
+        >
+          <TodoList />
+        </TodoProvider>
+      );
+
+      expect(getByText('help the town')).toBeInTheDocument();
+      expect(getByText('feed the dragon')).toBeInTheDocument();
+
+      act(() => getByText('remove help the town').click());
+
+      expect(queryByText(container, 'help the town')).not.toBeInTheDocument();
+      expect(getByText('feed the dragon')).toBeInTheDocument();
+    });
+  });
 });
diff --git a/packages/jotai-x/src/createAtomStore.ts b/packages/jotai-x/src/createAtomStore.ts
index fa675c8..8f5f769 100644
--- a/packages/jotai-x/src/createAtomStore.ts
+++ b/packages/jotai-x/src/createAtomStore.ts
@@ -1,17 +1,13 @@
-import { atom, useAtom, useAtomValue, useSetAtom } from 'jotai';
+import { useAtom, useAtomValue, useSetAtom } from 'jotai';
+import { useHydrateAtoms } from 'jotai/utils';
 
+import { atomWithFn } from './atomWithFn';
 import { createAtomProvider, useAtomStore } from './createAtomProvider';
 
 import type { ProviderProps } from './createAtomProvider';
 import type { FC } from 'react';
-import type { PrimitiveAtom } from 'jotai';
-import type { useHydrateAtoms } from 'jotai/utils';
-import type { createStore } from 'jotai/vanilla';
+import type { Atom, createStore, WritableAtom } from 'jotai/vanilla';
 
-type WithInitialValue<Value> = {
-  init: Value;
-};
-type Atom<Value> = PrimitiveAtom<Value> & WithInitialValue<Value>;
 export type JotaiStore = ReturnType<typeof createStore>;
 
 export type UseAtomOptions = {
@@ -22,17 +18,42 @@ export type UseAtomOptions = {
 
 type UseAtomOptionsOrScope = UseAtomOptions | string;
 
-export type GetRecord<O> = {
-  [K in keyof O]: (options?: UseAtomOptionsOrScope) => O[K];
+type GetRecord<O> = {
+  [K in keyof O]: O[K] extends Atom<infer V>
+    ? (options?: UseAtomOptionsOrScope) => V
+    : never;
+};
+
+type SetRecord<O> = {
+  [K in keyof O]: O[K] extends WritableAtom<infer _V, infer A, infer R>
+    ? (options?: UseAtomOptionsOrScope) => (...args: A) => R
+    : never;
+};
+
+type UseRecord<O> = {
+  [K in keyof O]: O[K] extends WritableAtom<infer V, infer A, infer R>
+    ? (options?: UseAtomOptionsOrScope) => [V, (...args: A) => R]
+    : never;
 };
-export type SetRecord<O> = {
-  [K in keyof O]: (options?: UseAtomOptionsOrScope) => (value: O[K]) => void;
+
+type StoreAtomsWithoutExtend<T> = {
+  [K in keyof T]: T[K] extends Atom<any> ? T[K] : SimpleWritableAtom<T[K]>;
 };
-export type UseRecord<O> = {
-  [K in keyof O]: (
-    options?: UseAtomOptionsOrScope
-  ) => [O[K], (value: O[K]) => void];
+
+type StoreAtoms<T, E> = StoreAtomsWithoutExtend<T> & E;
+
+type FilterWritableAtoms<T> = {
+  [K in keyof T]-?: T[K] extends WritableAtom<any, any, any> ? T[K] : never;
 };
+
+type WritableStoreAtoms<T, E> = FilterWritableAtoms<StoreAtoms<T, E>>;
+
+export type SimpleWritableAtom<T> = WritableAtom<T, [T], void>;
+
+export type SimpleWritableAtomRecord<T> = {
+  [K in keyof T]: SimpleWritableAtom<T[K]>;
+};
+
 export type AtomRecord<O> = {
   [K in keyof O]: Atom<O[K]>;
 };
@@ -51,32 +72,46 @@ export type UseSyncAtoms<T> = (
   }
 ) => void;
 
-export type StoreApi<T extends object, N extends string = ''> = {
-  atom: AtomRecord<T>;
+export type StoreApi<
+  T extends object,
+  E extends AtomRecord<object>,
+  N extends string = '',
+> = {
+  atom: StoreAtoms<T, E>;
   name: N;
-  extend: <ET, EN>(
-    extendedState: ET,
-    options?: Omit<
-      CreateAtomStoreOptions<object, EN extends string ? EN : N>,
-      'initialStore'
-    >
-  ) => AtomStoreApi<T & ET, EN extends string ? EN : N>;
 };
 
-export type UseStoreApi<T> = (options?: UseAtomOptionsOrScope) => {
-  get: GetRecord<T>;
-  set: SetRecord<T>;
-  use: UseRecord<T>;
+type GetAtomFn = <V>(atom: Atom<V>, options?: UseAtomOptionsOrScope) => V;
+
+type SetAtomFn = <V, A extends unknown[], R>(
+  atom: WritableAtom<V, A, R>,
+  options?: UseAtomOptionsOrScope
+) => (...args: A) => R;
+
+type UseAtomFn = <V, A extends unknown[], R>(
+  atom: WritableAtom<V, A, R>,
+  options?: UseAtomOptionsOrScope
+) => [V, (...args: A) => R];
+
+export type UseStoreApi<T, E> = (options?: UseAtomOptionsOrScope) => {
+  get: GetRecord<StoreAtoms<T, E>> & { atom: GetAtomFn };
+  set: SetRecord<WritableStoreAtoms<T, E>> & { atom: SetAtomFn };
+  use: UseRecord<WritableStoreAtoms<T, E>> & { atom: UseAtomFn };
+  store: (options?: UseAtomOptionsOrScope) => JotaiStore | undefined;
 };
 
-export type AtomStoreApi<T extends object, N extends string = ''> = {
+export type AtomStoreApi<
+  T extends object,
+  E extends AtomRecord<object>,
+  N extends string = '',
+> = {
   name: N;
 } & {
   [key in keyof Record<NameProvider<N>, object>]: FC<ProviderProps<T>>;
 } & {
-  [key in keyof Record<NameStore<N>, object>]: StoreApi<T, N>;
+  [key in keyof Record<NameStore<N>, object>]: StoreApi<T, E, N>;
 } & {
-  [key in keyof Record<UseNameStore<N>, object>]: UseStoreApi<T>;
+  [key in keyof Record<UseNameStore<N>, object>]: UseStoreApi<T, E>;
 };
 
 const capitalizeFirstLetter = (str = '') =>
@@ -88,12 +123,18 @@ const getStoreIndex = (name = '') =>
 const getUseStoreIndex = (name = '') =>
   `use${capitalizeFirstLetter(name)}Store`;
 
-const withDefaultOptions = <T>(
-  atomRecord: AtomRecord<T>,
+const isAtom = (possibleAtom: unknown): boolean =>
+  !!possibleAtom &&
+  typeof possibleAtom === 'object' &&
+  'read' in possibleAtom &&
+  typeof possibleAtom.read === 'function';
+
+const withDefaultOptions = <T extends object>(
+  fnRecord: T,
   defaultOptions: UseAtomOptions
-): AtomRecord<T> =>
+): T =>
   Object.fromEntries(
-    Object.entries(atomRecord).map(([key, fn]) => [
+    Object.entries(fnRecord).map(([key, fn]) => [
       key,
       (options: UseAtomOptions = {}) =>
         (fn as any)({ ...defaultOptions, ...options }),
@@ -107,12 +148,15 @@ const convertScopeShorthand = (
     ? { scope: optionsOrScope }
     : optionsOrScope;
 
-export interface CreateAtomStoreOptions<T extends object, N extends string> {
-  store?: UseAtomOptions['store'];
+export interface CreateAtomStoreOptions<
+  T extends object,
+  E extends AtomRecord<object>,
+  N extends string,
+> {
+  name: N;
   delay?: UseAtomOptions['delay'];
-  initialStore?: AtomStoreApi<T, N>;
-  name?: N;
   effect?: FC;
+  extend?: (atomsWithoutExtend: StoreAtomsWithoutExtend<T>) => E;
 }
 
 /**
@@ -129,94 +173,153 @@ export interface CreateAtomStoreOptions<T extends object, N extends string> {
  */
 export const createAtomStore = <
   T extends object,
-  IT extends object,
+  E extends AtomRecord<object>,
   N extends string = '',
 >(
   initialState: T,
-  {
-    delay: delayRoot,
-    initialStore,
-    name = '' as any,
-    effect,
-  }: CreateAtomStoreOptions<IT, N> = {}
-): AtomStoreApi<T & IT, N> => {
-  const useInitialStoreIndex = getUseStoreIndex(
-    initialStore?.name
-  ) as UseNameStore<N>;
-  const initialStoreIndex = getStoreIndex(initialStore?.name) as NameStore<N>;
+  { name, delay: delayRoot, effect, extend }: CreateAtomStoreOptions<T, E, N>
+): AtomStoreApi<T, E, N> => {
+  type MyStoreAtoms = StoreAtoms<T, E>;
+  type MyWritableStoreAtoms = WritableStoreAtoms<T, E>;
+  type MyStoreAtomsWithoutExtend = StoreAtomsWithoutExtend<T>;
+  type MyWritableStoreAtomsWithoutExtend =
+    FilterWritableAtoms<MyStoreAtomsWithoutExtend>;
+
   const providerIndex = getProviderIndex(name) as NameProvider<N>;
   const useStoreIndex = getUseStoreIndex(name) as UseNameStore<N>;
   const storeIndex = getStoreIndex(name) as NameStore<N>;
 
-  // FIXME: These constants have type any
-  const getAtoms = initialStore
-    ? (initialStore[useInitialStoreIndex] as any)().get
-    : ({} as GetRecord<T & IT>);
-  const setAtoms = initialStore
-    ? (initialStore[useInitialStoreIndex] as any)().set
-    : ({} as SetRecord<T & IT>);
-  const useAtoms = initialStore
-    ? (initialStore[useInitialStoreIndex] as any)().use
-    : ({} as UseRecord<T & IT>);
-  const atoms = initialStore
-    ? (initialStore[initialStoreIndex] as any).atom
-    : ({} as AtomRecord<T & IT>);
-
-  for (const key of Object.keys(initialState)) {
-    const atomConfig = atom(initialState[key as keyof T]);
-
-    atoms[key] = atomConfig;
-    getAtoms[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
-      const options = convertScopeShorthand(optionsOrScope);
-      const contextStore = useAtomStore(name, options.scope, false);
-
-      return useAtomValue(atomConfig, {
-        store: options.store ?? contextStore,
-        delay: options.delay ?? delayRoot,
-      });
-    };
-    setAtoms[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
-      const options = convertScopeShorthand(optionsOrScope);
-      const contextStore = useAtomStore(name, options.scope);
-
-      return useSetAtom(atomConfig as any, {
-        store: options.store ?? contextStore,
-      });
-    };
-    useAtoms[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
-      const options = convertScopeShorthand(optionsOrScope);
-      const contextStore = useAtomStore(name, options.scope);
-
-      return useAtom(atomConfig, {
-        store: options.store ?? contextStore,
-        delay: options.delay ?? delayRoot,
-      });
-    };
+  const atomsWithoutExtend = {} as MyStoreAtomsWithoutExtend;
+  const writableAtomsWithoutExtend = {} as MyWritableStoreAtomsWithoutExtend;
+  const atomIsWritable = {} as Record<keyof MyStoreAtoms, boolean>;
+
+  for (const [key, atomOrValue] of Object.entries(initialState)) {
+    const atomConfig: Atom<unknown> = isAtom(atomOrValue)
+      ? atomOrValue
+      : atomWithFn(atomOrValue);
+    atomsWithoutExtend[key as keyof MyStoreAtomsWithoutExtend] =
+      atomConfig as any;
+
+    const writable = 'write' in atomConfig;
+    atomIsWritable[key as keyof MyStoreAtoms] = writable;
+
+    if (writable) {
+      writableAtomsWithoutExtend[
+        key as keyof MyWritableStoreAtomsWithoutExtend
+      ] = atomConfig as any;
+    }
   }
 
-  const api: any = {
-    [providerIndex]: createAtomProvider(name, atoms, { effect }),
-    [useStoreIndex]: (options: UseAtomOptionsOrScope = {}) => ({
-      get: withDefaultOptions(getAtoms, convertScopeShorthand(options)),
-      set: withDefaultOptions(setAtoms, convertScopeShorthand(options)),
-      use: withDefaultOptions(useAtoms, convertScopeShorthand(options)),
-    }),
-    [storeIndex]: {
-      atom: atoms,
-      name,
-    },
+  const atoms = { ...atomsWithoutExtend } as MyStoreAtoms;
+
+  if (extend) {
+    const extendedAtoms = extend(atomsWithoutExtend);
+
+    for (const [key, atomConfig] of Object.entries(extendedAtoms)) {
+      atoms[key as keyof MyStoreAtoms] = atomConfig;
+      atomIsWritable[key as keyof MyStoreAtoms] = 'write' in atomConfig;
+    }
+  }
+
+  const getAtoms = {} as GetRecord<MyStoreAtoms>;
+  const setAtoms = {} as SetRecord<MyWritableStoreAtoms>;
+  const useAtoms = {} as UseRecord<MyWritableStoreAtoms>;
+
+  const useStore = (
+    optionsOrScope: UseAtomOptionsOrScope = {},
+    warnIfUndefined = true
+  ) => {
+    const { scope, store } = convertScopeShorthand(optionsOrScope);
+    const contextStore = useAtomStore(name, scope, warnIfUndefined);
+    return store ?? contextStore;
+  };
+
+  const useAtomValueWithStore: GetAtomFn = (atomConfig, optionsOrScope) => {
+    const store = useStore(optionsOrScope, false);
+    const { delay = delayRoot } = convertScopeShorthand(optionsOrScope);
+    return useAtomValue(atomConfig, { store, delay });
+  };
+
+  const useSetAtomWithStore: SetAtomFn = (atomConfig, optionsOrScope) => {
+    const store = useStore(optionsOrScope);
+    return useSetAtom(atomConfig, { store });
+  };
+
+  const useAtomWithStore: UseAtomFn = (atomConfig, optionsOrScope) => {
+    const store = useStore(optionsOrScope);
+    const { delay = delayRoot } = convertScopeShorthand(optionsOrScope);
+    return useAtom(atomConfig, { store, delay });
+  };
+
+  for (const key of Object.keys(atoms)) {
+    const atomConfig = atoms[key as keyof MyStoreAtoms];
+    const isWritable: boolean = atomIsWritable[key as keyof MyStoreAtoms];
+
+    (getAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
+      useAtomValueWithStore(atomConfig, optionsOrScope);
+
+    if (isWritable) {
+      (setAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
+        useSetAtomWithStore(
+          atomConfig as WritableAtom<any, any, any>,
+          optionsOrScope
+        );
+
+      (useAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
+        useAtomWithStore(
+          atomConfig as WritableAtom<any, any, any>,
+          optionsOrScope
+        );
+    }
+  }
+
+  const Provider: FC<ProviderProps<T>> = createAtomProvider(
+    name,
+    writableAtomsWithoutExtend,
+    { effect }
+  );
+
+  const storeApi: StoreApi<T, E, N> = {
+    atom: atoms,
     name,
   };
 
-  return {
-    ...api,
-    [storeIndex]: {
-      ...api[storeIndex],
-      extend: (extendedState: any, options: any) =>
-        createAtomStore(extendedState, {
-          initialStore: api,
-          ...options,
+  const useStoreApi: UseStoreApi<T, E> = (defaultOptions = {}) => ({
+    get: {
+      ...withDefaultOptions(getAtoms, convertScopeShorthand(defaultOptions)),
+      atom: (atomConfig, options) =>
+        useAtomValueWithStore(atomConfig, {
+          ...convertScopeShorthand(defaultOptions),
+          ...convertScopeShorthand(options),
         }),
     },
-  };
+    set: {
+      ...withDefaultOptions(setAtoms, convertScopeShorthand(defaultOptions)),
+      atom: (atomConfig, options) =>
+        useSetAtomWithStore(atomConfig, {
+          ...convertScopeShorthand(defaultOptions),
+          ...convertScopeShorthand(options),
+        }),
+    },
+    use: {
+      ...withDefaultOptions(useAtoms, convertScopeShorthand(defaultOptions)),
+      atom: (atomConfig, options) =>
+        useAtomWithStore(atomConfig, {
+          ...convertScopeShorthand(defaultOptions),
+          ...convertScopeShorthand(options),
+        }),
+    },
+    store: (options) =>
+      useStore({
+        ...convertScopeShorthand(defaultOptions),
+        ...convertScopeShorthand(options),
+      }),
+  });
+
+  return {
+    [providerIndex]: Provider,
+    [useStoreIndex]: useStoreApi,
+    [storeIndex]: storeApi,
+    name,
+  } as any;
 };
diff --git a/packages/jotai-x/src/index.ts b/packages/jotai-x/src/index.ts
index f63fe17..b9d1112 100644
--- a/packages/jotai-x/src/index.ts
+++ b/packages/jotai-x/src/index.ts
@@ -3,6 +3,7 @@
  */
 
 export * from './atomProvider';
+export * from './atomWithFn';
 export * from './createAtomProvider';
 export * from './createAtomStore';
 export * from './useHydrateStore';
diff --git a/packages/jotai-x/src/useHydrateStore.ts b/packages/jotai-x/src/useHydrateStore.ts
index 9dbc595..8af7ffe 100644
--- a/packages/jotai-x/src/useHydrateStore.ts
+++ b/packages/jotai-x/src/useHydrateStore.ts
@@ -2,8 +2,8 @@ import { useEffect } from 'react';
 import { useSetAtom } from 'jotai';
 import { useHydrateAtoms } from 'jotai/utils';
 
-import type {
-  AtomRecord,
+import {
+  SimpleWritableAtomRecord,
   UseHydrateAtoms,
   UseSyncAtoms,
 } from './createAtomStore';
@@ -12,7 +12,7 @@ import type {
  * Hydrate atoms with initial values for SSR.
  */
 export const useHydrateStore = (
-  atoms: AtomRecord<any>,
+  atoms: SimpleWritableAtomRecord<any>,
   initialValues: Parameters<UseHydrateAtoms<any>>[0],
   options: Parameters<UseHydrateAtoms<any>>[1] = {}
 ) => {
@@ -22,12 +22,7 @@ export const useHydrateStore = (
     const initialValue = initialValues[key];
 
     if (initialValue !== undefined) {
-      values.push([
-        atoms[key],
-        typeof initialValue === 'function'
-          ? { fn: initialValue }
-          : initialValue,
-      ]);
+      values.push([atoms[key], initialValue]);
     }
   }
 
@@ -38,7 +33,7 @@ export const useHydrateStore = (
  * Update atoms with new values on changes.
  */
 export const useSyncStore = (
-  atoms: AtomRecord<any>,
+  atoms: SimpleWritableAtomRecord<any>,
   values: any,
   { store }: Parameters<UseSyncAtoms<any>>[1] = {}
 ) => {
@@ -50,7 +45,7 @@ export const useSyncStore = (
 
     useEffect(() => {
       if (value !== undefined && value !== null) {
-        set(typeof value === 'function' ? { fn: value } : value);
+        set(value);
       }
     }, [set, value]);
   }