jupyter-notebook - Wrap text inside align-environment

Question

The following is a snippet from a jupyter-notebooks' markdown cell.

$$
\begin{aligned}
\mathrm{d}J & = \mathrm{tr}(\frac{\mathrm{d}J}{\mathrm{d}\mathbf{A}^{[3]}} \frac{\mathrm{d}\mathbf{A}^{[3]}}{\mathrm{d}\mathbf{Z}^{[3]}} \mathrm{d}\mathbf{Z}^{[3]})\\
&\begin{aligned}
& \text{Here, $\frac{\mathrm{d}\mathbf{A}^{[3]}}{\mathrm{d}\mathbf{Z}^{[3]}}$ should be a rank-4 tensor. But, the derivatives corresponding to a single training-example, i.e. $\frac{\mathrm{d}\mathbf{A}^{[3](i)}}{\mathrm{d}\mathbf{Z}^{[3](i)}}$, is a rank-2 tensor (more specifically, a 2D diagnonal-matrix), computed as follows: } \\
\frac{\mathrm{d}\mathbf{A}^{[3](i)}}{\mathrm{d}\mathbf{Z}^{[3](i)}} & = 
\biggl[\frac{\partial}{\partial{z}^{[3](i)}_1}, \frac{\partial}{\partial{z}^{[3](i)}_2}, \dots, \frac{\partial}{\partial{z}^{[3](i)}_{n^{[3]}}}\biggr]\otimes
\begin{bmatrix}
{a}^{[3](i)}_1\\
{a}^{[3](i)}_2\\
\vdots\\
{a}^{[3](i)}_{n^{[3]}}\\
\end{bmatrix} \\
& = \begin{bmatrix}
\frac{\partial {a}^{[3](i)}_1}{\partial{z}^{[3](i)}_1} & 0 & \dots & 0\\
0 & \frac{\partial {a}^{[3](i)}_2}{\partial{z}^{[3](i)}_2} & \dots & 0\\
\vdots & \vdots & \ddots & \vdots \\
0 & 0 & \dots & \frac{\partial {a}^{[3](i)}_{n^{[3]}}}{\partial{z}^{[3](i)}_{n^{[3]}}}\\
\end{bmatrix}
\end{aligned}
\end{aligned}
$$

Which renders as follows:

How do I make the text inside the \text{} to wrap?

PS: There are other posts (see below) that ask the same question. However, the solutions they offer (like using \parbox, \minipage, \intertext, etc.) are not working within Jupyter-notebook's markdown cell (which uses MathJax - see here).

• Wrap text around align* environment
• Plain text inside align* environment
• Wrap text inside align environment