Skip to content

Update Jax code#244

Open
cl126162 wants to merge 1 commit intomainfrom
cl/JAX_update
Open

Update Jax code#244
cl126162 wants to merge 1 commit intomainfrom
cl/JAX_update

Conversation

@cl126162
Copy link
Copy Markdown
Collaborator

@cl126162 cl126162 commented Apr 7, 2026

Kolmogorov flow:

  • refactor obs extracttion avoiding multiple FFT computations
  • add full jit support through entire env
  • FP32 support: stable at dt=7.5e-4 | FP64 support: stable at dt=1e-3

Channel flow:

  • full fp32 support
  • add full jit support through entire env

@smokbel Can you please carefully double-check if physical correctness is preserved?

@ludgerpaehler
Copy link
Copy Markdown
Collaborator

Any chance we can at least separate out the Nek components from this PR? Also, can we try to move the .yml configurations out into the HuggingFace repository?

@cl126162
Copy link
Copy Markdown
Collaborator Author

Configs are pushed now to HF

@cl126162
Copy link
Copy Markdown
Collaborator Author

nek and maia updates are now separated

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants