updated test app

app.py

@@ -1,5 +1,28 @@
 import streamlit as st
 
+import functools
+import e3x
+from flax import linen as nn
+import jax
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import numpy as np
+import optax
+
+# Disable future warnings.
+import warnings
+warnings.simplefilter(action='ignore', category=FutureWarning)
+
+# Initialize PRNGKey for random number generation.
+key = jax.random.PRNGKey(0)
+
+key, rotation_key = jax.random.split(key)
+rotation = e3x.so3.random_rotation(rotation_key)
+
+st.write(rotation)
+
+
+
 x = st.slider('Select a value')
 st.write(x, 'squared is', x * x)