Are custom PyTree classes supposed to be able to handle ShapedArray as input?

[dime10]

The question relates to the following scenario:

args_flat, args_tree = tree_flatten(args)
avals_flat = [shaped_abstractify(arg) for arg in args_flat]
tree_unflatten(avals_flat, args_tree)

Is this a pattern that would ever be used by JAX? In that case the constructor of a custom PyTree class which accepts arrays as inputs would also need to be able to handle ShapedArray inputs (let's assume it handles concrete arrays and tracers of abstract arrays just fine).

A relevant passage in this docs can be found here: https://jax.readthedocs.io/en/latest/pytrees.html#custom-pytrees-and-initialization, which warns against certain initialization code in custom PyTree classes:

For this reason, the init and new methods of custom PyTree classes should generally avoid doing any array conversion or other input validation, or else anticipate and handle these special cases.

But it doesn't mention the ShapedArray type explicitly or includes it in the example code for special handling:

class MyTree:
  def __init__(self, a):
    if not (type(a) is object or a is None or isinstance(a, MyTree)):
      a = jnp.asarray(a)
    self.a = a

[jakevdp]

I can't think of any situation where ShapedArray objects would be passed to your pytree in the course of normal execution. Generally ShapedArray values appear in the aval attribute of arrays and tracers, and it's these arrays and tracers that will be passed to your pytree constructor.

[dime10]

Thank you for your answer, that is helpful to know!

While searching I did manage to find one instance of this pattern in the code base:
https://github.com/google/jax/blob/9752ba9751e9bfb30b4006d94f4e7e82ba8f18f5/jax/_src/maps.py#L611-L620

I'm not sure about the context however.