Debugging

Debugging mode can be turned on globally by setting the FEEDBAX_DEBUG environment variable to "True".

This can be done from within Python, as shown below, or by setting FEEDBAX_DEBUG=True in your shell before executing Python from the command line.

import os

os.environ["FEEDBAX_DEBUG"] = str(True)

Logging details of model execution

The main debugging feature implemented in Feedbax so far is automatic logging of the state operations performed by any subclass of AbstractStagedModel.

For example, now that we've enabled debugging, when a model is executed during training or validation...

import jax

from feedbax.xabdeef import point_mass_nn_simple_reaches


key_init, key_eval = jax.random.split(jax.random.PRNGKey(0))
context = point_mass_nn_simple_reaches(key=key_init)
task, model = context.task, context.model

# We won't bother to train the model for this example

states = task.eval(model, key=key_eval)

...a detailed summary of the model stages that were executed will be saved to the log.

!cat feedbax.log

2024-03-12 12:19:12,023 [INFO] feedbax,63: Logger configured.
2024-03-12 12:19:15,694 [DEBUG] feedbax._staged,193: 
  Model type: Channel
  Stage: "update_queue"
  Callable:
      BoundMethod(
        __func__=<function _update_queue>,
        __self__=Channel(
          delay=1,
          noise_std=0.0,
          input_proto=(f32[2], f32[2]),
          intervenors={'update_queue': [], 'add_noise': []},
          _init_value=0.0
        )
      )
  Input:
      (f32[2], f32[2])
  Substate:
      ChannelState(output=(f32[2], f32[2]), queue=((f32[2], f32[2]),), noise=None)

2024-03-12 12:19:15,707 [DEBUG] feedbax._staged,193: 
  Model type: Channel
  Stage: "add_noise"
  Callable:
      BoundMethod(
        __func__=<function _add_noise>,
        __self__=Channel(
          delay=1,
          noise_std=0.0,
          input_proto=(f32[2], f32[2]),
          intervenors={'update_queue': [], 'add_noise': []},
          _init_value=0.0
        )
      )
  Input:
      (f32[2], f32[2])
  Substate:
      ((f32[2], f32[2]), (f32[2], f32[2]))

2024-03-12 12:19:15,710 [DEBUG] feedbax._staged,193: 
  Model type: SimpleFeedback
  Stage: "update_feedback"
  Callable:
      MultiModel(
        models=Channel(
          delay=1,
          noise_std=0.0,
          input_proto=(f32[2], f32[2]),
          intervenors={'update_queue': [], 'add_noise': []},
          _init_value=0.0
        )
      )
  Input:
      (f32[2], f32[2])
  Substate:
      ChannelState(
        output=(f32[2], f32[2]),
        queue=((f32[2], f32[2]),),
        noise=(f32[2], f32[2])
      )

2024-03-12 12:19:15,730 [DEBUG] feedbax._staged,193: 
  Model type: SimpleStagedNetwork
  Stage: "hidden"
  Callable:
      GRUCell(
        weight_ih=f32[150,8],
        weight_hh=f32[150,50],
        bias=f32[150],
        bias_n=f32[50],
        input_size=8,
        hidden_size=50,
        use_bias=True
      )
  Input:
      f32[8]
  Substate:
      f32[50]

2024-03-12 12:19:15,732 [DEBUG] feedbax._staged,193: 
  Model type: SimpleStagedNetwork
  Stage: "hidden_nonlinearity"
  Callable:
      <wrapped function identity_func>
  Input:
      f32[50]
  Substate:
      f32[50]

2024-03-12 12:19:15,736 [DEBUG] feedbax._staged,193: 
  Model type: SimpleStagedNetwork
  Stage: "readout"
  Callable:
      Linear(
        weight=f32[2,50],
        bias=f32[2],
        in_features=50,
        out_features=2,
        use_bias=True
      )
  Input:
      f32[50]
  Substate:
      f32[2]

2024-03-12 12:19:15,738 [DEBUG] feedbax._staged,193: 
  Model type: SimpleStagedNetwork
  Stage: "out_nonlinearity"
  Callable:
      <wrapped function identity_func>
  Input:
      f32[2]
  Substate:
      f32[2]

2024-03-12 12:19:15,741 [DEBUG] feedbax._staged,193: 
  Model type: SimpleFeedback
  Stage: "nn_step"
  Callable:
      SimpleStagedNetwork(
        hidden=GRUCell(
          weight_ih=f32[150,8],
          weight_hh=f32[150,50],
          bias=f32[150],
          bias_n=f32[50],
          input_size=8,
          hidden_size=50,
          use_bias=True
        ),
        hidden_size=50,
        hidden_noise_std=None,
        hidden_nonlinearity=<function identity_func>,
        out_size=2,
        out_nonlinearity=<function identity_func>,
        readout=Linear(
          weight=f32[2,50],
          bias=f32[2],
          in_features=50,
          out_features=2,
          use_bias=True
        ),
        encoding_size=None,
        encoder=None,
        intervenors={
          'hidden':
          [],
          'hidden_nonlinearity':
          [],
          'readout':
          [],
          'out_nonlinearity':
          []
        }
      )
  Input:
      (CartesianState(pos=f32[2], vel=f32[2], force=None), (f32[2], f32[2]))
  Substate:
      NetworkState(hidden=f32[50], output=f32[2], encoding=None)

2024-03-12 12:19:15,747 [DEBUG] feedbax._staged,193: 
  Model type: Mechanics
  Stage: "convert_effector_force"
  Callable:
      BoundMethod(
        __func__=<function update_state_given_effector_force>,
        __self__=PointMass(mass=1.0)
      )
  Input:
      f32[2]
  Substate:
      CartesianState(pos=f32[2], vel=f32[2], force=f32[2])

2024-03-12 12:19:15,751 [DEBUG] feedbax._staged,193: 
  Model type: DirectForceInput
  Stage: "clip_skeleton_state"
  Callable:
      BoundMethod(
        __func__=<function _clip_state>,
        __self__=DirectForceInput(
          skeleton=PointMass(mass=1.0),
          clip_states=True,
          intervenors={'clip_skeleton_state': []}
        )
      )
  Input:
      StateBounds(low=None, high=None)
  Substate:
      CartesianState(pos=f32[2], vel=f32[2], force=f32[2])

2024-03-12 12:19:15,752 [DEBUG] feedbax._staged,193: 
  Model type: Mechanics
  Stage: "kinematics_update"
  Callable:
      DirectForceInput(
        skeleton=PointMass(mass=1.0),
        clip_states=True,
        intervenors={'clip_skeleton_state': []}
      )
  Input:
      f32[2]
  Substate:
      PlantState(
        skeleton=CartesianState(pos=f32[2], vel=f32[2], force=f32[2]),
        muscles=None
      )

2024-03-12 12:19:15,782 [DEBUG] feedbax._staged,193: 
  Model type: Mechanics
  Stage: "dynamics_step"
  Callable:
      BoundMethod(
        __func__=<function dynamics_step>,
        __self__=Mechanics(
          plant=DirectForceInput(
            skeleton=PointMass(mass=1.0),
            clip_states=True,
            intervenors={'clip_skeleton_state': []}
          ),
          dt=0.05,
          solver=Euler(),
          intervenors={
            'convert_effector_force':
            [],
            'kinematics_update':
            [],
            'dynamics_step':
            [],
            'get_effector':
            []
          }
        )
      )
  Input:
      f32[2]
  Substate:
      MechanicsState(
        plant=PlantState(
          skeleton=CartesianState(pos=f32[2], vel=f32[2], force=f32[2]),
          muscles=None
        ),
        effector=CartesianState(pos=f32[2], vel=f32[2], force=f32[2]),
        solver=None
      )

2024-03-12 12:19:15,784 [DEBUG] feedbax._staged,193: 
  Model type: Mechanics
  Stage: "get_effector"
  Callable:
      BoundMethod(__func__=<wrapped function effector>, __self__=PointMass(mass=1.0))
  Input:
      CartesianState(pos=f32[2], vel=f32[2], force=f32[2])
  Substate:
      CartesianState(pos=f32[2], vel=f32[2], force=f32[2])

2024-03-12 12:19:15,786 [DEBUG] feedbax._staged,193: 
  Model type: SimpleFeedback
  Stage: "mechanics_step"
  Callable:
      Mechanics(
        plant=DirectForceInput(
          skeleton=PointMass(mass=1.0),
          clip_states=True,
          intervenors={'clip_skeleton_state': []}
        ),
        dt=0.05,
        solver=Euler(),
        intervenors={
          'convert_effector_force':
          [],
          'kinematics_update':
          [],
          'dynamics_step':
          [],
          'get_effector':
          []
        }
      )
  Input:
      f32[2]
  Substate:
      MechanicsState(
        plant=PlantState(
          skeleton=CartesianState(pos=f32[2], vel=f32[2], force=f32[2]),
          muscles=None
        ),
        effector=CartesianState(pos=f32[2], vel=f32[2], force=f32[2]),
        solver=None
      )

Each entry in the log begins with the date and time the log was made.

Entries whose first line contains [DEBUG] feedbax.staged are logs of model stages that were just executed. Thus, the time between the previous log entry and the current one gives the approximate duration of the current stage.

Following that first line are indented lines that provide details:

Model typeStageCallableInputSubstate

The type of AbstractStagedModel that the stage that was just executed belongs to.

The stage's label, which is used to identify it among all the stages that belong to its model type.

What is actually called, in order to execute the stage. This may be:

• an Equinox model—such as another type of AbstractStagedModel—in which case the code that was executed can be found in that module's __call__ method. For example, in the stage "hidden" of model type SimpleStagedNetwork logged above, the Equinox module GRUCell was called;
• a method of a class or module. Typically these are logged as BoundMethod, where __func__ provides the name of the method, and __self__ provides the class that the method belongs to;
• some other function or callable object.

Logs the structure of the input that was passed to the callable.

The part of the model type's state that can be updated by the model stage.

For example, the state associated with the staged model SimpleFeedback is SimpleFeedbackState, which contains a field network of type NetworkState. The stage "nn_step" of SimpleFeedback passes this NetworkState to SimpleStagedNetwork which returns an updated NetworkState.

Warning

Keep in mind that for long training runs, enabling FEEDBAX_DEBUG may result in very large log files, with all of the model stages logged repeatedly on every training run. If you are trying to debug a structural problem in the model step, it's preferable to do a single evaluation of the model as above, so that each stage will only be executed once.

Debugging in JAX

JAX provides several debugging features which you should be familiar with while developing Feedbax models.