JAX training with marimo
The JAX tutorial deploys a JAX example notebook to CKS with a single CLI command. The notebook trains a small identity network using Flax and Optax, and streams a live loss chart to the browser as training progresses. The example uses the standardghcr.io/marimo-team/marimo:latest image, no custom Dockerfile needed. Dependencies including jax[cuda12] are pinned in a PEP 723 header embedded in the notebook file, so the environment is recreated exactly on every pod start.
TensorRT-LLM inference with marimo
The TensorRT-LLM tutorial runs NVIDIA TensorRT-LLM inference inside a marimo notebook on CKS using the official NVIDIA Triton NGC container. The notebook includes an interactive model picker and prompt selector. Supported models include TinyLlama 1.1B, Phi-3.5-mini, Mistral 7B, Llama-3.1 8B FP8, and Minitron 8B. The tutorial supports both thekubectl-marimo CLI plugin and a plain YAML manifest.
Both tutorials require the marimo operator installed on your cluster.