Remove references to deprecated jax.nn.normalize

jax.nn.standardize is a drop-in replacement after deprecation (see jax-ml/jax#18439) PiperOrigin-RevId: 580728570
google · Nov 9, 2023 · ac28613 · ac28613
1 parent 50de4c4
commit ac28613
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/flax/linen/activation.py b/flax/linen/activation.py
@@ -33,7 +33,6 @@
 from jax.nn import log_sigmoid
 from jax.nn import log_softmax
 from jax.nn import logsumexp
-from jax.nn import normalize
 from jax.nn import one_hot
 from jax.nn import relu
 from jax.nn import relu6
@@ -48,6 +47,9 @@
 import jax.numpy as jnp
 from jax.numpy import tanh
 
+# Normalize is a deprecated alias of standardize
+normalize = standardize
+
 # pylint: enable=unused-import