Feature/jax simplify visualization#184

Jammy2211 merged 2 commits intofeature/jax_wrapperJammy2211/PyAutoArray:feature/jax_wrapperfrom

feature/jax_simplify_visualizationJammy2211/PyAutoArray:feature/jax_simplify_visualizationCopy head branch name to clipboard