Add CuTe DSL JAX demo

[katjasrz]

Summary

This PR adds a minimal CuTe DSL + JAX demo under examples/python/CuTeDSL/jax.

The demo shows how to define and invoke custom CuTe DSL kernels from JAX, providing a self-contained example that can serve as a starting point for experimentation and integration.

What This Adds

• cute_dsl_jax.ipynb
  A walkthrough-style notebook demonstrating:

  • Defining CuTe DSL kernels
  • Calling them from JAX
  • Basic usage patterns and expected behavior

• cute_dsl_jax_kernels.py
  Supporting kernel definitions used by the notebook.

Purpose

The goal is to provide:

• A simple reference example for users exploring CuTe DSL + JAX
• A starting template for extending CuTe DSL kernels within a JAX workflow
• A lightweight demo that complements the existing Python examples

Notes

• The example is self-contained and lives under examples/python/CuTeDSL/jax.
• No changes to core functionality.
• Intended primarily as an illustrative example.

[brandon-yujie-sun]

Great, thanks for adding the example and notebook!

[brandon-yujie-sun]

Would it make more sense to move this nb file to examples/python/CuTeDSL/notebooks, which would make the testing easier to have all nb files in consolidated place?

[katjasrz]

As agreed on another thread, let's keep the current location