Skip to main content
Ctrl+K
This site covers the new Flax NNX API. [Click here for the old Flax Linen API]
 - Home
  • Flax basics
  • MNIST tutorial
  • Why Flax NNX?
  • Guides
    • Transformations
    • Scale up on multiple devices
    • Using Filters, grouping NNX variables
    • Randomness
    • Performance considerations
    • Evolution from Flax Linen to NNX
    • Use Flax NNX and Linen together
    • Model surgery
    • Save and load checkpoints
    • Flax NNX vs JAX transformations
    • Migrating from Haiku to Flax
    • Example: Using pretrained Gemma for inference with Flax NNX
  • Examples
    • Core examples
  • Flax NNX glossary
  • The Flax philosophy
  • How to contribute
  • API Reference
    • flax.nnx
      • graph
      • module
      • nn
        • Activation functions
        • Attention
        • Dtypes
        • Initializers
        • Linear
        • LoRA
        • Normalization
        • Recurrent
        • Stochastic
      • rnglib
      • spmd
      • state
      • training
        • Metrics
        • Optimizer
      • transforms
      • variables
      • helpers
      • visualization
      • filterlib
      • bridge
    • flax.core.frozen_dict package
    • flax.struct package
    • flax.training package
  • .rst

filterlib

Contents

  • to_predicate()
  • WithTag
  • PathContains
  • OfType
  • Any
  • All
  • Not
  • Everything
  • Nothing

filterlib#

flax.nnx.filterlib.to_predicate(filter)[source]#

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

class flax.nnx.WithTag(tag: str)[source]#
class flax.nnx.PathContains(key: flax.typing.Key)[source]#
class flax.nnx.OfType(type: type)[source]#
class flax.nnx.Any(*filters)[source]#
class flax.nnx.All(*filters)[source]#
class flax.nnx.Not(collection_filter, /)[source]#
class flax.nnx.Everything[source]#
class flax.nnx.Nothing[source]#

previous

visualization

next

bridge

Contents
  • to_predicate()
  • WithTag
  • PathContains
  • OfType
  • Any
  • All
  • Not
  • Everything
  • Nothing

By The Flax authors

© Copyright 2023, The Flax authors.