Vectorization and promotion #13149

bkaplowitz · 2022-11-07T19:23:37Z

bkaplowitz
Nov 7, 2022

Hi,
When I run my code I get the following error message.

Traceback (most recent call last):
  File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/api_util.py", line 395, in flatten_axes
    tree_map(add_leaves, _replace_nones(proxy, axis_tree), dummy)
  File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/tree_util.py", line 204, in tree_map
    all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
  File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/tree_util.py", line 204, in <listcomp>
    all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
ValueError: Tuple arity mismatch: 2 != 1; tuple: (FrozenDict({
    params: {
        hidden: {
            bias: <object object at 0x7fd6d808ae10>,
            kernel: <object object at 0x7fd6d808ae10>,
        },
        hidden_1: {
            hf: {
                bias: <object object at 0x7fd6d808ae10>,
                kernel: <object object at 0x7fd6d808ae10>,
            },
            hg: {
                bias: <object object at 0x7fd6d808ae10>,
                kernel: <object object at 0x7fd6d808ae10>,
            },
            hi: {
                bias: <object object at 0x7fd6d808ae10>,
                kernel: <object object at 0x7fd6d808ae10>,
            },
            ho: {
                bias: <object object at 0x7fd6d808ae10>,
                kernel: <object object at 0x7fd6d808ae10>,
            },
            if: {
                kernel: <object object at 0x7fd6d808ae10>,
            },
            ig: {
                kernel: <object object at 0x7fd6d808ae10>,
            },
            ii: {
                kernel: <object object at 0x7fd6d808ae10>,
            },
            io: {
                kernel: <object object at 0x7fd6d808ae10>,
            },
        },
        logits: {
            bias: <object object at 0x7fd6d808ae10>,
            kernel: <object object at 0x7fd6d808ae10>,
        },
        value: {
            bias: <object object at 0x7fd6d808ae10>,
            kernel: <object object at 0x7fd6d808ae10>,
        },
    },
}), <object object at 0x7fd6d808ae10>).

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/train.py", line 75, in <module>
    algo.train(rng)
  File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 187, in train
    val, metric, traj_batch = self._update_step(val, None)
  File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 195, in _update_step
    new_train_state, info = self._update(val[0], traj_batch)
  File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 280, in _update
    val, info = jax.lax.scan(self._update_epoch, val, None, self.config["NUM_EPOCHS_PER_UPDATE"])
  File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/lax/control_flow/loops.py", line 231, in scan
    carry, y = f(carry, tree_unflatten(xs_tree, xs_slice))
  File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 305, in _update_epoch
    total_loss, grads = batched_grad_fn(
  File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/api.py", line 1562, in vmap_f
    in_axes_flat = flatten_axes("vmap in_axes", in_tree, (in_axes, 0), kws=True)
  File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/api_util.py", line 401, in flatten_axes
    assert treedef_is_leaf(leaf)
jax._src.traceback_util.UnfilteredStackTrace: AssertionError

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 187, in train
    val, metric, traj_batch = self._update_step(val, None)
  File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 195, in _update_step
    new_train_state, info = self._update(val[0], traj_batch)
  File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 280, in _update
    val, info = jax.lax.scan(self._update_epoch, val, None, self.config["NUM_EPOCHS_PER_UPDATE"])
  File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 305, in _update_epoch
    total_loss, grads = batched_grad_fn(
AssertionError

I'm not exactly sure what is going on so any help would be appreciated. In particular, I don't know what the arity mismatch means.

Second, and perhaps relatedly, sometimes in vmapped functions when running over a batch, python sometimes promotes to (1,1,2,4) and sometimes to (1,2,4). I've tried tracing back the source of this consistency with no success, especially as once I check for it and everything passes I run into the arity error.

The relevant portion of the function where the error occurs is given by the last line of the following code:

    def _update_epoch(self, val, unused):
        batch, gae, targets, train_state = val
        jax.debug.print("{jax_val}", jax_val=jax.value_and_grad(loss_actor_and_critic, has_aux=False))
        print(batch.obs.shape, "batch obs shape")
        batched_grad_fn = jax.vmap(
            jax.value_and_grad(loss_actor_and_critic, has_aux=False),  # True, included extra info
            in_axes=(
                (
                    0,
                    None,
                    None,
                    None
                ),
            ),
        )
        # axis_name = ("batch",)
        # $), in_axes=(None,0,  None), axis_name="agents",
        # )
        action = batch.action[:-1]


        total_loss, grads = batched_grad_fn(
            train_state.params,
            self.network.apply,
            state=batch.obs[:-1],
            target=targets,
            value_old=batch.value[:-1],
            log_pi_old=batch.log_prob[:-1],
            gae=gae,
            action=action,
            clip_eps=self.config["eps_clip"],
            critic_coeff=self.config["critic_coef"],
            entropy_coeff=self.config["entropy_coef"],
            carry=(batch.carry0[:-1], batch.carry1[:-1]),
        )

        train_state = train_state.apply_gradients(grads=grads)
        return (batch, targets, gae, train_state), total_loss

Finally, the function with jax value and grad is given by:

def loss_actor_and_critic(
    params_model: flax.core.frozen_dict.FrozenDict,
    apply_fn: Callable[..., Any],
    state: jnp.ndarray,
    target: jnp.ndarray,
    value_old: jnp.ndarray,
    log_pi_old: jnp.ndarray,
    gae: jnp.ndarray,
    action: jnp.ndarray,
    clip_eps: float,
    critic_coeff: float,
    entropy_coeff: float,
    carry: Tuple[jnp.ndarray, jnp.ndarray],
) -> jnp.ndarray:
    # jax.debug.breakpoint()

    carry, value_pred, pi = apply_fn(params_model, state, carry)
    # reshaped_action = jnp.expand_dims(action)
    log_prob = jnp.expand_dims(pi.log_prob(action), -1)
    value_pred_clipped = value_old + (value_pred - value_old).clip(-clip_eps, clip_eps)
    value_losses = jnp.square(value_pred - target)
    value_losses_clipped = jnp.square(value_pred_clipped - target)
    value_loss = 0.5 * jnp.maximum(value_losses, value_losses_clipped).mean(axis=0)

    ratio = jnp.exp(log_prob - log_pi_old)
    # compute surrogate loss function
    gae = (gae - gae.mean()) / (gae.std() + 1e-8)
    # Compute the policy loss with clipped surrogate objective.
    loss_actor1 = ratio * gae
    loss_actor2 = jnp.clip(ratio, 1.0 - clip_eps, 1.0 + clip_eps) * gae
    loss_actor = -jnp.minimum(loss_actor1, loss_actor2)
    loss_actor = loss_actor.mean()
    # Compute entropy bonus.
    entropy = pi.entropy().mean()
    # Compute the total loss as the sum of the actor and critic losses, minus an entropy bonus.
    total_loss: jnp.ndarray = loss_actor + critic_coeff * value_loss - entropy_coeff * entropy
    # Optional returns: (value_loss, loss_actor, entropy, value_pred.mean(), target.mean(), gae.mean())
    return total_loss.squeeze()

Any suggestions would be greatly appreciated.

jakevdp · 2022-11-07T19:42:22Z

jakevdp
Nov 7, 2022
Maintainer

Hi - this is a very confusing error message, but the issue comes from the fact that your in_axes specification doesn't match the arguments you're using to call the function. Your in-axes specifies a single argument that is a tuple of four arrays; you call the function with two positional arguments and a dozen keyword arguments. This mismatch in length of arguments is what is causing the error.

It's hard to tell exactly what to recommend as a fix without being able to run the code and know a bit more about your intent, but I would start there.

3 replies

bkaplowitz Nov 7, 2022
Author

This is a good catch. Unfortunately changing to in_axes = (0,None), changed nothing about the error as far as I can tell. I can attach runnable code if it would be helpful.

Traceback (most recent call last):
File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/api_util.py", line 395, in flatten_axes
tree_map(add_leaves, _replace_nones(proxy, axis_tree), dummy)
File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/tree_util.py", line 204, in tree_map
all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/tree_util.py", line 204, in
all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
ValueError: Tuple arity mismatch: 2 != 1; tuple: (FrozenDict({
params: {
hidden: {
bias: <object object at 0x7f97c0f96e10>,
kernel: <object object at 0x7f97c0f96e10>,
},
hidden_1: {
hf: {
bias: <object object at 0x7f97c0f96e10>,
kernel: <object object at 0x7f97c0f96e10>,
},
hg: {
bias: <object object at 0x7f97c0f96e10>,
kernel: <object object at 0x7f97c0f96e10>,
},
hi: {
bias: <object object at 0x7f97c0f96e10>,
kernel: <object object at 0x7f97c0f96e10>,
},
ho: {
bias: <object object at 0x7f97c0f96e10>,
kernel: <object object at 0x7f97c0f96e10>,
},
if: {
kernel: <object object at 0x7f97c0f96e10>,
},
ig: {
kernel: <object object at 0x7f97c0f96e10>,
},
ii: {
kernel: <object object at 0x7f97c0f96e10>,
},
io: {
kernel: <object object at 0x7f97c0f96e10>,
},
},
logits: {
bias: <object object at 0x7f97c0f96e10>,
kernel: <object object at 0x7f97c0f96e10>,
},
value: {
bias: <object object at 0x7f97c0f96e10>,
kernel: <object object at 0x7f97c0f96e10>,
},
},
}), <object object at 0x7f97c0f96e10>).

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/train.py", line 75, in
algo.train(rng)
File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 187, in train
val, metric, traj_batch = self._update_step(val, None)
File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 195, in _update_step
new_train_state, info = self._update(val[0], traj_batch)
File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 280, in _update
val, info = jax.lax.scan(self._update_epoch, val, None, self.config["NUM_EPOCHS_PER_UPDATE"])
File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/lax/control_flow/loops.py", line 231, in scan
carry, y = f(carry, tree_unflatten(xs_tree, xs_slice))
File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 303, in _update_epoch
total_loss, grads = batched_grad_fn(
File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/api.py", line 1562, in vmap_f
in_axes_flat = flatten_axes("vmap in_axes", in_tree, (in_axes, 0), kws=True)
File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/site-packages/jax/_src/api_util.py", line 401, in flatten_axes
assert treedef_is_leaf(leaf)
jax._src.traceback_util.UnfilteredStackTrace: AssertionError

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 187, in train
val, metric, traj_batch = self._update_step(val, None)
File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 195, in _update_step
new_train_state, info = self._update(val[0], traj_batch)
File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 280, in _update
val, info = jax.lax.scan(self._update_epoch, val, None, self.config["NUM_EPOCHS_PER_UPDATE"])
File "/Users/brandonkaplowitz/Documents/Documents - Comp/Graduate Research/My Projects/aiyagari_christian_brandon_jax/src/ppo.py", line 303, in _update_epoch
total_loss, grads = batched_grad_fn(
AssertionError
Exception ignored in: <module 'threading' from '/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/threading.py'>
Traceback (most recent call last):
File "/Users/brandonkaplowitz/opt/anaconda3/envs/aiyagari_christian_brandon_jax/lib/python3.10/threading.py", line 1560, in _shutdown
lock.acquire()
KeyboardInterrupt:

jakevdp Nov 7, 2022
Maintainer

I'm still having trouble understanding your intent here. Note that in vmapped fuctions:

Arguments passed as keywords are always mapped over their leading axis (i.e. axis index 0).

(See https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html)

Is this the behavior you want in the keyword arguments to your vmapped function? From the names of the parameters, I suspect some of your keyword arguments are not arrays.

bkaplowitz Nov 10, 2022
Author

Oh, I see. No this is not what I want. The types are:

params_model: flax.core.frozen_dict.FrozenDict,
apply_fn: Callable[..., Any],
state: jnp.ndarray,
target: jnp.ndarray,
value_old: jnp.ndarray,
log_pi_old: jnp.ndarray,
gae: jnp.ndarray,
action: jnp.ndarray,
clip_eps: float,
critic_coeff: float,
entropy_coeff: float,
carry: Tuple[jnp.ndarray, jnp.ndarray],

I want the ndarray's to be mapped over, but not the config terms that are type of float, or the apply_fn.

bkaplowitz · 2022-11-15T14:14:45Z

bkaplowitz
Nov 15, 2022
Author

I have two quick follow-up questions when you have a chance.

Keyword args are always mapped over as you said, so does this mean that I have to have a wrapper function to load the keywords and then vmap the interior function? This seems somewhat inelegant, but if it is the only way to handle keywords, it makes sense. I guess 'None' is ignored for keywords/they aren't counted as part of the possible in_axes? I almost wonder then if it would be possible to somehow indicate auxiliary arguments. I'll test this out!
Is there an advantage of implementing grad_and_value surrounded by a vmap so that the entries are 1d arrays vs. the batched_grad implementation discussed in the documentation (https://jax.readthedocs.io/en/latest/jax-101/04-advanced-autodiff.html) or the reduce_axes keyword arg for grad?

Thanks so much for all of the help! I really appreciate the advice that you gave so far.

1 reply

jakevdp Nov 17, 2022
Maintainer

I don't think there's any way currently for vmap to have unmapped auxilary keyword arguments. Closure is the way to do this.

Regarding vmap of value_and_grad, this can be useful if you want to apply gradients element-wise.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Vectorization and promotion #13149

Uh oh!

{{title}}

Uh oh!

Replies: 2 comments 4 replies

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{editor}}'s edit

{{editor}}'s edit

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{editor}}'s edit

{{editor}}'s edit

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Select a reply

Uh oh!

Vectorization and promotion #13149

Uh oh!

bkaplowitz Nov 7, 2022

Replies: 2 comments · 4 replies

Uh oh!

Uh oh!

jakevdp Nov 7, 2022 Maintainer

Uh oh!

bkaplowitz Nov 7, 2022 Author

Uh oh!

Uh oh!

jakevdp Nov 7, 2022 Maintainer

Uh oh!

bkaplowitz Nov 10, 2022 Author

Uh oh!

bkaplowitz Nov 15, 2022 Author

Uh oh!

jakevdp Nov 17, 2022 Maintainer

bkaplowitz
Nov 7, 2022

Replies: 2 comments 4 replies

jakevdp
Nov 7, 2022
Maintainer

bkaplowitz Nov 7, 2022
Author

jakevdp Nov 7, 2022
Maintainer

bkaplowitz Nov 10, 2022
Author

bkaplowitz
Nov 15, 2022
Author

jakevdp Nov 17, 2022
Maintainer