remove inlined jax.nn.initializers definitions, resolving TODO of lev…

…skaya et al fixes breakage from cl/655766534 aka #21069 PiperOrigin-RevId: 655783604
jax-ml · Jul 25, 2024 · a94aea5 · a94aea5
1 parent 76b4c70
commit a94aea5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/jax/_src/api.py b/jax/_src/api.py
@@ -2693,9 +2693,9 @@ class ShapeDtypeStruct:
     sharding: (optional) a :class:`jax.Sharding` object
   """
   __slots__ = ["shape", "dtype", "sharding", "_dll"]
-  named_shape = {}
+  named_shape = {}  # type: ignore
 
-  def __init__(self, shape, dtype, sharding=None, named_shape=None):
+  def __init__(self, shape, dtype, named_shape=None, sharding=None):
     del named_shape  # ignored, vestigial
     self.shape = tuple(shape)
     if dtype is None: