filterlib

flax.nnx.filterlib.to_predicate(filter)

Converts a Filter to a predicate function. See Using Filters.

class flax.nnx.WithTag(tag: str)

class flax.nnx.PathContains(key: flax.typing.Key | str, exact: bool = True)

class flax.nnx.OfType(type: type)

class flax.nnx.Any(*filters)

class flax.nnx.All(*filters)

class flax.nnx.Not(collection_filter, /)

class flax.nnx.Everything

class flax.nnx.Nothing