GitHub repository: ARHashemi/Neural_Net_JAX_Optax_Keras
Description: A simple neural network is created using JAX with optax optimizer and a custom defined loss function.