jax: TypeError when using None as custom_vjp cotangent for custom pytree and asarray

import jax.numpy as jnp
from jax.tree_util import register_pytree_node_class

@register_pytree_node_class
class C:
    def __init__(self, a=1.):
        self.a = jnp.asarray(a, dtype=float)  # TypeError: float() argument must be a string or a real number, not 'object'
        # self.a = jnp.asarray(a)  # TypeError: Value '<object object at 0x7fff5220fc20>' with dtype object is not a valid JAX array type. Only arrays of numeric types are supported by JAX.

    def tree_flatten(self):
        return (self.a,), None

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        return cls(*children)

@custom_vjp
def f(x, y):
    return 2. * x  # doesn't depend on y in this simple example


def f_fwd(x, y):
    z = f(x, y)
    res = None
    return z, res

def f_bwd(res, z_cot):
    x_cot = 2. * z_cot
    y_cot = None
    return x_cot, y_cot

f.defvjp(f_fwd, f_bwd)

c = C()
vjp(f, 3., c)[1](3.)  # TypeError with different messages (see above) depending on whether c.a is weakly typed or not
vjp(f, 3., [1., 2])[1](3.)  # this works