Add jax.tree.static() - SemanticDiff

Commit

98 days ago

Add jax.tree.static() This makes it more convenient to specify static arguments with jax.tree_util.register_dataclass.