Remove deprecated function as_named_shape.

PiperOrigin-RevId: 655920752
google-research · Jul 25, 2024 · e2e9312 · e2e9312
1 parent ab86199
commit e2e9312
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/e3x/nn/initializers.py b/e3x/nn/initializers.py
@@ -178,7 +178,7 @@ def _random_array(
       return stddev * _complex_truncated_normal(
           key=key,
           upper=2,
-          shape=jax.core.as_named_shape(shape),
+          shape=shape,
           dtype=dtype,
       )
   elif distribution == 'normal':
@@ -192,7 +192,7 @@ def _random_array(
       )
     else:
       return jnp.sqrt(variance) * _complex_uniform(
-          key=key, shape=jax.core.as_named_shape(shape), dtype=dtype
+          key=key, shape=shape, dtype=dtype
       )
   else:
     raise ValueError(f"invalid distribution '{distribution}' for _random_array")