flax
permit passing mesh or rules to spmd primitives
#2939

Merged

Commits

permit passing mesh or rules to spmd primitives

cgarciae committed 3 years ago
add rules and mesh docstrings

cgarciae committed 3 years ago
ignore FutureWarning from jax

cgarciae committed 3 years ago
fix mypy errors

cgarciae committed 3 years ago

Loading