filterlib

flax.nnx.filterlib.to_predicate(filter)

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

class flax.nnx.WithTag(tag)

Filter that matches values with a string tag attribute equal to the given tag.

Used by RngKey and RngCount. See Using Filters.

tag

The string tag to match against a value’s tag attribute.

Type:

str

class flax.nnx.PathContains(key, exact=True)

Filter that matches values whose associated path contains the given key.

See Using Filters.

key

The key to look for in the value’s path.

Type:

flax.typing.Key | str

exact

If True, key must equal a path element exactly. If False, matches when key is contained in the string form of any path element.

Type:

bool

class flax.nnx.OfType(type)

Filter that matches values that are instances of type, or that have a type attribute that is an instance of type.

See Using Filters.

type

The type to match values against.

Type:

type

class flax.nnx.Any(*filters)

Filter that matches values matching any of the inner filters.

See Using Filters.

class flax.nnx.All(*filters)

Filter that matches values matching all of the inner filters.

See Using Filters.

class flax.nnx.Not(collection_filter, /)

Filter that matches values that do not match the inner filter.

See Using Filters.

class flax.nnx.Everything

Filter that matches all values.

See Using Filters.

class flax.nnx.Nothing

Filter that matches no values.

See Using Filters.