Function count_model_params

axlearn/common/utils.py:1331–1333

Count the number of parameters in a model.

(tree: NestedTensor)

        cast_fn = per_param_dtype_by_path(default_dtype, update_rules)
        per_param_dtype = cast_fn(tree)
        Result:
        per_param_dtype = {
            'conv1_weights': jnp.float32,
            'conv2_weights': jnp.float32,
            'fc1_weights': jnp.bfloat16,
            'fc2_weights': jnp.bfloat16,