PyPI page
Home page
Author:
Ashish Kumar Singh
Summary:
A JAX/Flax framework for training diffusion, JEPA and other large models from scratch
Latest version:
0.1.0
Required dependencies:
einops
|
filelock
|
flax
|
grain
|
huggingface-hub
|
jax
|
jinja2
|
numpy
|
opencv-python-headless
|
optax
|
orbax-checkpoint
|
pillow
|
rich
|
transformers
|
typing-extensions
|
tyro
Optional dependencies:
compressed-tensors
|
cosmic-ray
|
datasets
|
diffusers
|
flatbuffers
|
ftfy
|
gguf
|
ipykernel
|
jax-cuda12-plugin
|
jax-cuda13-plugin
|
libtpu
|
lm-eval
|
matplotlib
|
mlflow
|
moviepy
|
nbclient
|
nbformat
|
ollama
|
openai
|
optuna
|
outlines-core
|
pytest
|
pytest-xdist
|
qwix
|
requests
|
ruff
|
safetensors
|
scipy
|
starlette
|
tensorboard
|
tensorflow-datasets
|
torch
|
torchax
|
torchvision
|
transformers
|
uvicorn
|
wandb
|
xprof
Downloads last day:
3
Downloads last week:
180
Downloads last month:
0
7-day smoothing