app.py

import streamlit as st
import jax

x = st.slider('Select a value')
st.write(x, 'squared is', x * x)

st.write(jax.device_count())