FROM docker.io/python:3.12.7-bookworm

WORKDIR /workspace

RUN apt-get update && apt-get install -y \
    build-essential \
    git \
    curl \
    && rm -rf /var/lib/apt/lists/*

RUN pip install --no-cache-dir \
    jax[cuda12] \
    equinox \
    torch \
    optax \
    torchvision

CMD ["/usr/bin/bash"]