Patrick Kidger

Probably something like this: ```python from typing import Union, TYPE_CHECKING if TYPE_CHECKING: from torch import TorchTensor from numpy import ndarray from jaxtyping import Array as JaxArray from tensorflow import TfTensor...

Actually, the more general `Exception` is deliberate. There are cases when try-importing a module can result in other issues too, c.f. https://github.com/google/jaxtyping/blob/7a84b27da9e57c425ce4e6333121c3cdf2e07302/jaxtyping/_array_types.py#L33-L39 As for adding the above to jaxtyping. jaxtyping...

I think what you're measuring here is an additive overhead of microseconds in the flattening and unflattening. This is expected/known, and rarely ever troublesome. Your computational work has to be...

@johannahaffner thank you for identifying that this is still the additive overhead! @nstarman I think what'd probably be most desirable is if we can just speed up the existing flattening/unflattening...

What happens if you just add the decorator on top of an existing class? How does the dataclass approach benchmark against codegen approach above? I'd like to avoid adding extra...

> Is there any reason Equinox couldn't do the exact same using its own flatten/unflatten!? I don't think this works with custom `__init__` functions: https://github.com/jax-ml/jax/issues/25486 Although if you have a...

I don't think that works, for the same reason as the issue I link: ```python import dataclasses import jax._src.tree_util @dataclasses.dataclass class X: x: int def __init__(self, x: str): assert type(x)...

I suspect option 2 might be reasonable. FWIW JAX already doesn't support the full gamut of dataclass things -- in particular as here, custom `__init__` methods -- so it's already...

> On a separate note: I think that a potential speedup to initialization could happen by avoiding all the _make_initable_wrapper(cls) machinery if there wasn't a custom __init__. On this, I...

This runs afoul of the 'no custom attributes on `Module`' rule we have. Ever-so-hypothetically a user may have defined a `_HAS_DATACLASS_INIT` attribute themselves. It's a bit of a philosophical point...