penzai icon indicating copy to clipboard operation
penzai copied to clipboard
verified publisher icon google-deepmind

Fix dtype type annotations to use `jax.typing.DTypeLike`.

Open Jammf opened this issue 3 months ago • 0 comments

The named_axis functions full(), zeros() and ones() annotate their dtype parameters with the non-existent np.DTypeLike. This change replaces the annotations with jax.typing.DTypeLike | None to match the dtype parameter of their wrapped functions.

Jammf avatar Oct 16 '25 20:10 Jammf