Search results

[Docs] JAX_Vision_transformer: ViT Tutorial at https://docs.jaxstack.ai/en/latest/JAX_Vision_transformer.html Fails on Kaggle (TPU v5e-8) or T4 (Colab)

Status: Open.

#281 In jax-ml/jax-ai-stack;

· Prakharprasun opened on Dec 4, 2025

Running the JAX_for_LLM_pretraining.ipynb in colab fails

Status: Open.

#265 In jax-ml/jax-ai-stack;

· Davidnet opened on Nov 10, 2025

Segmentation fault when initializing UNet model - (from tutorial problem)

Status: Open.

#261 In jax-ml/jax-ai-stack;

· Tenchberry opened on Oct 31, 2025

in digits_vae.ipynb, error: no attribute ModelAndOptimizer

Status: Open.

#241 In jax-ml/jax-ai-stack;

· joshtburdick opened on Oct 6, 2025

Execution errors in "Porting a PyTorch model to JAX"

Status: Open.

#218 In jax-ml/jax-ai-stack;

· pavithraes opened on Jun 18, 2025

Image captioning tutorial - <code>transformers</code> library with JAX

Status: Open.

#214 In jax-ml/jax-ai-stack;

· 8bitmp3 opened on Jun 12, 2025

Add LoRA Training Example to JAX Examples

Status: Open.

#186 In jax-ml/jax-ai-stack;

· nikolasavic3 opened on Mar 21, 2025

Discussion : What is the functionality of <code>pygrain.ShardOption</code>, like <code>pygrain.ShardByJaxProcess</code>?

Status: Open.

#154 In jax-ml/jax-ai-stack;

· carlesoctav opened on Jan 22, 2025

Colab XLA Error in JAX_for_LLM_pretraining.ipynb

Status: Open.

#149 In jax-ml/jax-ai-stack;

· Matt-Hurd opened on Jan 13, 2025