Below is the proof of equivalence between two different statements of Axiom of Choice
Axiom of Choice Statement 1 :
$f:X \mapsto Y$ is onto $\Rightarrow \exists g:Y \mapsto X$ s.t. $f\circ g= 1_Y$
Axiom of Choice Statment 2:
Let $\Bbb{P}$ be a partition of X. $\exists g:\Bbb{P} \mapsto X$ for every $A \in \Bbb{P}$ s.t. $ g(A) \in A$
Axiom of Choice Statment 3:
For given $X$, $\exists h: 2^X\setminus \{\emptyset\}\rightarrow X$ s.t. $h(A) \in A$ for each $A \in 2^X\setminus \{\emptyset\}$
Axiom of Choice Statment 4:
$\prod_{i\in I}X_i \neq \emptyset $ if $I\neq \emptyset$ and $X _i\neq$ for $ \forall i\in I$
Claim: Statement 1 $\Rightarrow$ Statement2
Let $\Bbb{P}$ be a partition of X. Then $A$ is uniquely given for each $x \in X$ s.t. $x \in A \in \Bbb P $
Let this unique-given be $f : X \mapsto \Bbb P$
This $f$ is $onto$ since every $A \in \Bbb P$ and there $\exists$ at least one $x \in X$ since $\Bbb P$ is partition
Now, from the Statement 1, since $f$ is $onto$, $\exists g:\Bbb P \mapsto X$ s.t. $f\circ g= 1_\Bbb P$
Since $f \circ g (A) = A$, $g(A) \in A $ thus,
$\exists g:\Bbb{P} \mapsto X $ for every $A \in \Bbb{P} $ s.t. $ g(A) \in A$
Claim: Statement 2 $\Rightarrow$ Statement1
$f:X \rightarrow Y$ is onto $\Rightarrow \Bbb P:= \{f^{-1}(\{y\})\mid y \in Y\}$ be a partition of $X$. Thus
$\exists g: \Bbb P \rightarrow X$ s.t. $g[\{f^{-1}(\{y\})] \in f^{-1}(\{y\})$
Define $h: Y \rightarrow X $ s.t. $h(y):=g[f^{-1}(\{y\})]$ then
$f(h(y))=y$ since $h(y) = f^{-1}(\{y\})$
Claim: Statement 2 $\Rightarrow$ Statement 3
Define $\widetilde{A} := \{(A,a) \mid a \in A\}\subset (2^X\setminus\{\emptyset\}) \times X$ then
$\Bbb P = \{\widetilde{A}\mid A\in 2^X\setminus\{\emptyset\}\}$ be a partition of
$Y =\cup\{\widetilde{A}\mid A\in 2^X\setminus\{\emptyset\}\}$ thus
$\exists g: \Bbb P \rightarrow Y s.t. g(\widetilde{A})\in \widetilde{A}$ for each $A \in 2^X\setminus\{\emptyset\}$ ( Statement 2)
Now there $\exists h: A\rightarrow \pi_2(g(\widetilde{A})), A\in2^X\setminus\{\emptyset\} $ s.t. $\pi_2 : (2^X\setminus\{\emptyset\})\times X \rightarrow X$
Claim: Statement 3 $\rightarrow$ Statement 2
It is true since a partition $\Bbb P \subset 2^X\setminus\{\emptyset\}$
Claim: Statement 3 $\rightarrow$ Statement 4
Let $X = \cup_{i\in I}X_i$ and $g: I \rightarrow X$ be $g(i)=h(X_i)$
then $g(i)=h(X_i) \in X$ for $h: 2^X\setminus \{\emptyset\}\rightarrow X$
Claim: Statement 4 $\rightarrow$ Statement 3
$\prod \{A\mid A\in 2^X\setminus \{\emptyset\}\} \neq \emptyset $
Let $h \in \prod \{A\mid A\in 2^X\setminus \{\emptyset\}\}$ then $h(A) \in A$
It concludes that abvoe 4 statments logically equivalent.
Any illogical part in this proof?
Your proof is mathematically sound.
The key point to notice here is that every partition of $X$ induces a set $Y$ and a surjection from $X$ onto $Y$; and every function with domain $X$ induces a partition of $X$ as well (and every function is onto its range). This is why these two are equivalent.
For a nice additional statement, show that every equivalence relation $E$ on a set $X$ has some $R\subseteq X$ such that every $x\in X$ is $E$-equivalent to exactly one element of $R$. (The key here is to notice that equivalence relations and partitions are "morally the same".)